import cv2
import numpy as np
import os

class FileStorage(object):
    def __init__(self, filename, isWrite=False):
        version = cv2.__version__
        self.major_version = int(version.split('.')[0])
        self.second_version = int(version.split('.')[1])

        if isWrite:
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_WRITE)
        else:
            self.fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_READ)

    def __del__(self):
        cv2.FileStorage.release(self.fs)

    def write(self, key, value, dt='mat'):
        if dt == 'mat':
            cv2.FileStorage.write(self.fs, key, value)
        elif dt == 'list':
            if self.major_version == 4: # 4.4
                self.fs.startWriteStruct(key, cv2.FileNode_SEQ)
                for elem in value:
                    self.fs.write('', elem)
                self.fs.endWriteStruct()
            else: # 3.4
                self.fs.write(key, '[')
                for elem in value:
                    self.fs.write('none', elem)
                self.fs.write('none', ']')

    def read(self, key, dt='mat'):
        if dt == 'mat':
            output = self.fs.getNode(key).mat()
        elif dt == 'list':
            results = []
            n = self.fs.getNode(key)
            for i in range(n.size()):
                val = n.at(i).string()
                if val == '':
                    val = str(int(n.at(i).real()))
                if val != 'none':
                    results.append(val)
            output = results
        else:
            raise NotImplementedError
        return output

    def close(self):
        self.__del__(self)

def read_intri(intri_name):
    assert os.path.exists(intri_name), intri_name
    intri = FileStorage(intri_name)
    camnames = intri.read('names', dt='list')
    cameras = {}
    for key in camnames:
        cam = {}
        cam['K'] = intri.read('K_{}'.format(key))
        cam['invK'] = np.linalg.inv(cam['K'])
        cam['dist'] = intri.read('dist_{}'.format(key))
        cameras[key] = cam
    return cameras

def write_intri(intri_name, cameras):
    intri = FileStorage(intri_name, True)
    results = {}
    camnames = list(cameras.keys())
    intri.write('names', camnames, 'list')
    for key_, val in cameras.items():
        key = key_.split('.')[0]
        K, dist = val['K'], val['dist']
        assert K.shape == (3, 3), K.shape
        assert dist.shape == (1, 5) or dist.shape == (5, 1), dist.shape
        intri.write('K_{}'.format(key), K)
        intri.write('dist_{}'.format(key), dist.reshape(1, 5))

def write_extri(extri_name, cameras):
    extri = FileStorage(extri_name, True)
    results = {}
    camnames = list(cameras.keys())
    extri.write('names', camnames, 'list')
    for key_, val in cameras.items():
        key = key_.split('.')[0]
        extri.write('R_{}'.format(key), val['Rvec'])
        extri.write('Rot_{}'.format(key), val['R'])
        extri.write('T_{}'.format(key), val['T'])
    return 0

def read_camera(intri_name, extri_name, cam_names=[]):
    assert os.path.exists(intri_name), intri_name
    assert os.path.exists(extri_name), extri_name

    intri = FileStorage(intri_name)
    extri = FileStorage(extri_name)
    cams, P = {}, {}
    cam_names = intri.read('names', dt='list')
    for cam in cam_names:
        # 内参只读子码流的
        cams[cam] = {}
        cams[cam]['K'] = intri.read('K_{}'.format( cam))
        cams[cam]['invK'] = np.linalg.inv(cams[cam]['K'])
        Rvec = extri.read('R_{}'.format(cam))
        Tvec = extri.read('T_{}'.format(cam))
        R = cv2.Rodrigues(Rvec)[0]
        RT = np.hstack((R, Tvec))

        cams[cam]['RT'] = RT
        cams[cam]['R'] = R
        cams[cam]['T'] = Tvec
        P[cam] = cams[cam]['K'] @ cams[cam]['RT']
        cams[cam]['P'] = P[cam]

        cams[cam]['dist'] = intri.read('dist_{}'.format(cam))
    cams['basenames'] = cam_names
    return cams

def write_camera(camera, path):
    from os.path import join
    intri_name = join(path, 'intri.yml')
    extri_name = join(path, 'extri.yml')
    intri = FileStorage(intri_name, True)
    extri = FileStorage(extri_name, True)
    results = {}
    camnames = [key_.split('.')[0] for key_ in camera.keys()]
    intri.write('names', camnames, 'list')
    extri.write('names', camnames, 'list')
    for key_, val in camera.items():
        if key_ == 'basenames':
            continue
        key = key_.split('.')[0]
        intri.write('K_{}'.format(key), val['K'])
        intri.write('dist_{}'.format(key), val['dist'])
        if 'Rvec' not in val.keys():
            val['Rvec'] = cv2.Rodrigues(val['R'])[0]
        extri.write('R_{}'.format(key), val['Rvec'])
        extri.write('Rot_{}'.format(key), val['R'])
        extri.write('T_{}'.format(key), val['T'])

def camera_from_img(img):
    height, width = img.shape[0], img.shape[1]
    # focal = 1.2*max(height, width) # as colmap
    focal = 1.2*min(height, width) # as colmap
    K = np.array([focal, 0., width/2, 0., focal, height/2, 0. ,0., 1.]).reshape(3, 3)
    camera = {'K':K ,'R': np.eye(3), 'T': np.zeros((3, 1)), 'dist': np.zeros((1, 5))}
    return camera

class Undistort:
    @staticmethod
    def image(frame, K, dist):
        return cv2.undistort(frame, K, dist, None)

    @staticmethod
    def points(keypoints, K, dist):
        # keypoints: (N, 3)
        assert len(keypoints.shape) == 2, keypoints.shape
        kpts = keypoints[:, None, :2]
        kpts = np.ascontiguousarray(kpts)
        kpts = cv2.undistortPoints(kpts, K, dist, P=K)
        keypoints[:, :2] = kpts[:, 0]
        return keypoints

    @staticmethod
    def bbox(bbox, K, dist):
        keypoints = np.array([[bbox[0], bbox[1], 1], [bbox[2], bbox[3], 1]])
        kpts = Undistort.points(keypoints, K, dist)
        bbox = np.array([kpts[0, 0], kpts[0, 1], kpts[1, 0], kpts[1, 1], bbox[4]])
        return bbox

def undistort(camera, frame=None, keypoints=None, output=None, bbox=None):
    # bbox: 1, 7
    print('This function is deprecated')
    raise NotImplementedError

def get_fundamental_matrix(cameras, basenames):
    skew_op = lambda x: np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
    fundamental_op = lambda K_0, R_0, T_0, K_1, R_1, T_1: np.linalg.inv(K_0).T @ (
            R_0 @ R_1.T) @ K_1.T @ skew_op(K_1 @ R_1 @ R_0.T @ (T_0 - R_0 @ R_1.T @ T_1))
    fundamental_RT_op = lambda K_0, RT_0, K_1, RT_1: fundamental_op (K_0, RT_0[:, :3], RT_0[:, 3], K_1,
                                                                          RT_1[:, :3], RT_1[:, 3] )
    F = np.zeros((len(basenames), len(basenames), 3, 3))  # N x N x 3 x 3 matrix
    F = {(icam, jcam): np.zeros((3, 3)) for jcam in basenames for icam in basenames}
    for icam in basenames:
        for jcam in basenames:
            F[(icam, jcam)] += fundamental_RT_op(cameras[icam]['K'], cameras[icam]['RT'], cameras[jcam]['K'], cameras[jcam]['RT'])
            if F[(icam, jcam)].sum() == 0:
                F[(icam, jcam)] += 1e-12  # to avoid nan
    return F


#!/usr/bin/env python
# -*-coding:utf-8 -*-
"""
  ████
██    ██   Datature
  ██  ██   Powering Breakthrough AI
    ██

@File    :   annotation_maker.py
@Author  :   Leonard So
@Version :   1.0
@Contact :   hello@datature.io
@License :   Apache License 2.0
@Desc    :   Automatic bounding box annotation creation for whole image.
"""
import os
import csv
import argparse
from PIL import Image

def create_annotation_file(image_dir):
    """
    Creates CSV Four Corner - Bounding Box annotation file to upload to
    the Nexus platform. This should be used on local storage of images.
    The file structure should be as follows (according to your current
    tags):
    ├── OK
    │   ├── img1.jpg
    │   ├── ...
    ├── Flashing
    │   ├── img2.jpg
    │   ├── ...
    ├── ShortMould
    │   ├── img3.jpg
    │   ├── ...
    ├── Flashing_ShortMould
    │   ├── img4.jpg
    │   ├── ...
    Requires Pillow to run, so please install in virtual environment.
    Args:
      image_dir: string for folder path containing the above folder
      structure. The corresponding csv file will be saved in the
      folder as annotations.csv. Once it has been created, you can
      upload the annotation with the format as CSV Four Corner -
      Bounding Box.
    """
    f = open(image_dir + "/annotations.csv", "w")
    writer = csv.writer(f)
    header = ["filename", "xmin", "ymin", "xmax", "ymax", "label"]
    writer.writerow(header)
    acceptable_ext = [".png", ".jpg", ".jpeg"]
    for subdir, _, files in os.walk(image_dir):
        label = subdir.split("/")[-1]
        for file in files:
            filename, ext = os.path.splitext(file)
            if ext in acceptable_ext:
                width, height = Image.open(os.path.join(subdir, file)).size
                writer.writerow([filename, 0, 0, width, height, label])

def args_parser():
    """
    Creates arg parser for command line input.
    """
    parser = argparse.ArgumentParser(
        description="Datature Open Source Prediction Script"
    )
    parser.add_argument(
        "--dataset",
        help="Path to folder that contains image dataset",
        required=True,
    )
    return parser.parse_args()

if __name__ == """__main__""":
    args = args_parser()
    if os.path.exists(args.dataset) is False:
        raise Exception("Folder Path Does Not Exist")
    create_annotation_file(args.dataset)


import ass
from subprocess import call
import datetime
import sys

episode_number = str(sys.argv[1])
ass_file = "subs-" + episode_number + ".ass"
mkv_file = "Twin.Peaks.S03E" + episode_number + ".mkv"

with open(ass_file, "r") as f:
    doc = ass.parse(f)
    x = 0
    while x < len(doc.events):
        time = doc.events[x].start + datetime.timedelta(milliseconds=250)
        seconds = time.total_seconds()
        call([
            "ffmpeg",
            "-ss",
            str(time),
            "-i",
            mkv_file,
            "-copyts",
            "-af",
            "asetpts=PTS-"+str(seconds)+"/TB",
            "-vf",
            "ass="+ass_file+",setpts=PTS-"+str(seconds)+"/TB",
            "-frames:v",
            "1",
            episode_number+"-"+str(x)+".png"
        ])
        x +=1


import plotly.express as px
import geopandas as gpd
import pandas as pd
from plot.colors.func import get_colors
def ie():
    colors=get_colors(r"source\plot\colors\no_yes\ny.txt")
    world=gpd.read_file("data\countries.geojson")
    data=pd.read_csv(r"data\budget-to-manage-invasive-alien-species.csv")
    data.rename(columns={"15.8.1 - Countries with an allocation from the national budget to manage the threat of invasive alien species (1 = YES, 0 = NO) - ER_IAS_NATBUD":"Budget for invasive species"},inplace=True)
    data["Budget for invasive species"].replace([0,1],["No","Yes"],inplace=True)
    full=world.merge(data,left_on=["ISO_A3"],right_on=["Code"])
    fig = px.choropleth(full,locations="ISO_A3",color="Budget for invasive species",hover_data=['ADMIN'],color_discrete_sequence=colors,width=1000)
    fig.update_layout(title="Countries with a budget for invasive alien species management, 2020")
    return fig


from typing import Tuple

import numpy as np

from ipie.trial_wavefunction.noci import NOCI
from ipie.trial_wavefunction.particle_hole import ParticleHole, ParticleHoleNonChunked
from ipie.trial_wavefunction.single_det import SingleDet
from ipie.trial_wavefunction.wavefunction_base import TrialWavefunctionBase
from ipie.utils.io import (
    determine_wavefunction_type,
    read_noci_wavefunction,
    read_particle_hole_wavefunction,
    read_qmcpack_wfn_hdf,
    read_single_det_wavefunction,
)

def get_trial_wavefunction(
    num_elec: Tuple[int, int],
    nbasis: int,
    wfn_file: str,
    ndets: int = -1,
    ndets_props: int = -1,
    ndet_chunks: int = 1,
    verbose=False,
):
    """Wavefunction factory.

    Parameters
    ----------
    options : dict
        Trial wavefunction input options.
    verbose : bool
        Print information.

    Returns
    -------
    trial : class or None
        Trial wavfunction class.
    """
    assert ndets_props <= ndets
    wfn_type = determine_wavefunction_type(wfn_file)
    if wfn_type == "particle_hole":
        wfn, _ = read_particle_hole_wavefunction(wfn_file)
        if ndet_chunks == 1:
            trial = ParticleHoleNonChunked(
                wfn,
                num_elec,
                nbasis,
                num_dets_for_trial=ndets,
                num_dets_for_props=ndets_props,
                verbose=verbose,
            )
        else:
            trial = ParticleHole(
                wfn,
                num_elec,
                nbasis,
                num_dets_for_trial=ndets,
                num_dets_for_props=ndets_props,
                num_det_chunks=ndet_chunks,
                verbose=verbose,
            )
    elif wfn_type == "noci":
        wfn, _ = read_noci_wavefunction(wfn_file)
        ci, (wfna, wfnb) = wfn
        assert len(wfn) == 2
        na = wfna.shape[-1]
        nb = wfnb.shape[-1]
        _nbasis = wfna.shape[0]
        assert nbasis == _nbasis
        outwfn = np.zeros((wfna.shape[0], wfna.shape[1], na + nb), dtype=wfna.dtype)
        outwfn[:, :, :na] = wfna.copy()
        outwfn[:, :, na:] = wfnb.copy()
        trial = NOCI((ci, outwfn), (na, nb), nbasis)
    elif wfn_type == "single_determinant":
        wfn, _ = read_single_det_wavefunction(wfn_file)
        assert len(wfn) == 2
        na = wfn[0].shape[-1]
        nb = wfn[1].shape[-1]
        _nbasis = wfn[0].shape[0]
        assert nbasis == _nbasis
        trial = SingleDet(np.hstack(wfn), (na, nb), nbasis)
    elif wfn_type == "qmcpack":
        trial = setup_qmcpack_wavefunction(wfn_file, ndets, ndets_props, ndet_chunks)
    else:
        raise RuntimeError("Unknown wavefunction type")
    trial.build()

    return trial

def setup_qmcpack_wavefunction(
    wfn_file: str, ndets: int, ndets_props: int, ndet_chunks: int
) -> TrialWavefunctionBase:
    wfn, psi0, nelec = read_qmcpack_wfn_hdf(wfn_file, get_nelec=True)
    nbasis = psi0.shape[0]
    if len(wfn) == 3:
        if ndet_chunks == 1:
            trial = ParticleHoleNonChunked(
                wfn, nelec, nbasis, num_dets_for_trial=ndets, num_dets_for_props=ndets_props
            )
        else:
            trial = ParticleHole(
                wfn,
                nelec,
                nbasis,
                num_dets_for_trial=ndets,
                num_dets_for_props=ndets_props,
                num_det_chunks=ndet_chunks,
            )
    elif len(wfn) == 2:
        if len(wfn[0]) == 1:
            nbasis = wfn[1][0].shape[0]
            trial = SingleDet(wfn[1][0], nelec, nbasis)
        else:
            nbasis = wfn[1][0].shape[0]
            trial = NOCI(wfn, nelec, nbasis)
    else:
        raise RuntimeError("Unknown QMCPACK wavefunction format.")
    return trial

def chunk_trial(self, handler):
    self.chunked = True  # Boolean to indicate that chunked cholesky is available

    if handler.scomm.rank == 0:  # Creating copy for every rank == 0
        self._rchola = self._rchola.copy()
        self._rcholb = self._rcholb.copy()

    self._rchola_chunk = handler.scatter_group(self._rchola)  # distribute over chol
    self._rcholb_chunk = handler.scatter_group(self._rcholb)  # distribute over chol

    tot_size = handler.allreduce_group(self._rchola_chunk.size)
    assert self._rchola.size == tot_size
    tot_size = handler.allreduce_group(self._rcholb_chunk.size)
    assert self._rcholb.size == tot_size


import re
import sys
from google.oauth2 import service_account
from google.auth.transport import requests
import base64
from pathlib import Path
import argparse
from typing import Optional
from typing import Sequence
import json
import os

API_BASE_URL='https://backstory.googleapis.com'
SCOPES = ['https://www.googleapis.com/auth/chronicle-backstory']

def get_http_client_from_file(service_account_file):
  # Create a credential using Google Developer Service Account Credential and Chronicle API scope.
  credentials = service_account.Credentials.from_service_account_file(service_account_file, scopes=SCOPES)
  # Build an HTTP client to make authorized OAuth requests.
  return requests.AuthorizedSession(credentials)

def get_http_client_from_sa_info(service_account_info):
  service_account_info = os.environ.get("SA_CREDENTIAL", "Not Found")
  # Create a credential using Google Developer Service Account Credential and Chronicle API scope.
  credentials = service_account.Credentials.from_service_account_info(json.loads(service_account_info), scopes=SCOPES)
  # Build an HTTP client to make authorized OAuth requests.
  return requests.AuthorizedSession(credentials)

def get_http_client_from_sa_env(env_var_name):
  service_account_info = os.environ.get(env_var_name)
  # Create a credential using Google Developer Service Account Credential and Chronicle API scope.
  credentials = service_account.Credentials.from_service_account_info(json.loads(service_account_info), scopes=SCOPES)
  # Build an HTTP client to make authorized OAuth requests.
  return requests.AuthorizedSession(credentials)

def initialize_command_line_args(
    args: Optional[Sequence[str]] = None) -> Optional[argparse.Namespace]:
  """Initializes and checks all the command-line arguments."""
  parser = argparse.ArgumentParser()
  parser.add_argument(
    "-c", "--credentials_file", type=str, help="path to credentials file")
  parser.add_argument(
    "-i", "--credentials_info", type=str, help="service account credentials info (alternative to file)")
  parser.add_argument(
    "-e", "--credentials_env", type=str, help="service account credentials info stored in environment variable")
  parser.add_argument(
    "-l", "--local_path", type=str, required=True, help="local rules path")
  parser.add_argument(
    "-m", "--make_changes", action='store_true',  help="fix any differences, if omitted then just report differences" )
  parser.add_argument(
    "-s", "--silent", action='store_true',  help="supress error messages" )
  parser.add_argument(
    "-r", "--region", required=False, default="us", choices=("asia-southeast1", "europe", "us"), help="Chronicle instance region (leave blank for US)" )

  parser.set_defaults(make_changes=False)
  return parser.parse_args(args)

def get_all_rules(session):
  try:
    rules_response = []
    url=f'{API_BASE_URL}/v2/detect/rules'
    response = session.request(
      method="GET",
      url=url,
    )
    if response.json():
      rules_array = response.json()
      for rule in rules_array['rules']:
        rules_response.append({
          'name': rule['ruleName'],
          'rule_id': rule['ruleId'],
          'text': base64.b64encode(rule['ruleText'].strip().encode('utf-8')).decode('utf-8'),
          'text_raw': rule['ruleText'].strip(),
        })
    return rules_response
  except:
    print(f'Received error code: {response.status_code}')

def load_local_files(path):
  rules_response = []
  file_list = list(Path(path).rglob("*.[yY][aA][rR][aA][lL]"))
  for file in file_list:
    file_content = open(file, "r").read()
    rules_response.append({
      'name': parse_rule_name(file_content), #str(file.stem),
      'path': str(file),
      'text': base64.b64encode(file_content.strip().encode('utf-8')).decode('utf-8'),
      'text_raw': file_content.strip(),
    })
  return rules_response

def parse_rule_name(rule_text):
  rule_name = re.search("rule ([A-Za-z0-9_]+)[\r\n\s]?{", rule_text)
  if rule_name:
    return rule_name.group(1)
  else:
    return False

def compare_rules(session, local_rules_path, silent=False):
  count_non_existent_rules = 0
  count_matched_rules = 0
  count_rules_to_update = 0
  local_rules = load_local_files(local_rules_path)
  remote_rules = get_all_rules(session)
  for local_rule in local_rules:
    if remote_rules:
      matched_remote_rules = [x for x in remote_rules if x['name'] == local_rule['name']]
      if len(matched_remote_rules) > 0:
        if len(matched_remote_rules) != 1:
          print(f'Matched rules was not equal to 1 for {local_rule["name"]}')
        else:
          if matched_remote_rules[0]['text'] == local_rule['text']:
            count_matched_rules += 1
          else:
            count_rules_to_update += 1
      else:
        count_non_existent_rules += 1
    else:
      count_non_existent_rules += 1
  if remote_rules:
    count_remote_rules = len(remote_rules)
  else:
    count_remote_rules = 0
  if local_rules:
    count_local_rules = len(local_rules)
  else:
    count_local_rules = 0
  return {
    "matched_rules": count_matched_rules,
    "rules_to_update": count_rules_to_update,
    "non_existent_rules": count_non_existent_rules,
    "remote_rules_total": count_remote_rules,
    "local_rules_total": count_local_rules,
  }

def upload_missing_rules(session, local_rules_path, silent=False):
  rules_updated = 0
  rules_added = 0
  error_count = 0
  errors = []
  local_rules = load_local_files(local_rules_path)
  remote_rules = get_all_rules(session)
  for local_rule in local_rules:
    if local_rule['name']:
      matched_remote_rules = [x for x in remote_rules if x['name'] == local_rule['name']]
      if len(matched_remote_rules) > 0:
        if len(matched_remote_rules) != 1:
          print(f'upload_missing_rules / Matched rules was not equal to 1 for {local_rule["name"]}')
        else:
          try:
            if matched_remote_rules[0]['text'] != local_rule['text']:
              if update_rule(session,local_rule['text_raw'],matched_remote_rules[0]['rule_id'],local_rule['name'], silent):
                rules_updated += 1
              else:
                error_count += 1
          except:
            if not silent: print(f"upload_missing_rules / Update of rule {local_rule['name']} failed")
            errors.append(f"upload_missing_rules / Update of rule {local_rule['name']} failed")
            error_count += 1
      else:
        try:
          if create_rule(session,local_rule['text_raw'],local_rule['name'],silent):
            rules_added += 1
          else:
            error_count += 1
        except:
          if not silent: print(f"upload_missing_rules / Creation of rule {local_rule['name']} failed")
          errors.append(f"upload_missing_rules / Creation of rule {local_rule['name']} failed")
          error_count += 1
    else:
      if not silent: print(f"upload_missing_rules / rule with path {local_rule['path']} has an issue")
      errors.append(f"upload_missing_rules / Creation of rule {local_rule['path']} failed")
      error_count += 1

  return {
    "rules_added": rules_added,
    "rules_updated": rules_updated,
    "errors": error_count,
    "error_details": errors,
  }

def verify_rule(session,rule_text,rule_name,silent=False):
  try:
    url=f'{API_BASE_URL}/v2/detect/rules:verifyRule'
    body={"ruleText":rule_text}
    response = session.request(
      method="POST",
      url=url,
      json=body
    )
    response.raise_for_status()
    return True
  except:
    if not silent: print(f'verify_rule [{rule_name}]/ Received error code: {response.status_code}')
    return False

def create_rule(session,rule_text,rule_name, silent=False):
  if verify_rule(session,rule_text,rule_name,silent):
    try:
      url=f'{API_BASE_URL}/v2/detect/rules'
      body={"ruleText":rule_text}
      response = session.request(
        method="POST",
        url=url,
        json=body
      )
      response.raise_for_status()
      return True
    except:
      if not silent: print(f'create_rule / Received error code: {response.status_code}')
      return False
  else:
    return False

def update_rule(session,rule_text,rule_id,rule_name,silent=False):
  if verify_rule(session,rule_text,rule_name,silent):
    try:
      url=f'{API_BASE_URL}/v2/detect/rules/{rule_id}:createVersion'
      body={"ruleText":rule_text}
      response = session.request(
        method="POST",
        url=url,
        json=body
      )
      response.raise_for_status()
      return True
    except:
      if not silent: print(f'update_rule / Received error code: {response.status_code}')
      return False
  else:
    return False

if __name__ == "__main__":
  cli = initialize_command_line_args()
  if not cli:
    sys.exit(1)  # A sanity check failed.
  if not cli.region == "us":
    API_BASE_URL=f'https://{cli.region}-backstory.googleapis.com'
  if cli.credentials_file:
    session = get_http_client_from_file(cli.credentials_file)
  elif cli.credentials_info:
    session = get_http_client_from_sa_info(cli.credentials_info)
  elif cli.credentials_env:
    session = get_http_client_from_sa_env(cli.credentials_env)
  else:
    print("No credential passed")
  if not cli.make_changes:
    print(json.dumps(compare_rules(session,cli.local_path, cli.silent),indent=2))
  else:
    print(json.dumps(upload_missing_rules(session,cli.local_path, cli.silent),indent=2))


import os
import asyncio

import cv2
import numpy as np
import torch
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from ultralytics import YOLO
from tqdm import tqdm

from misc import PlotPoseData

class DistProcessing:

    def __init__(self, yolo_model, folder_path: str, save_results: bool = True):
        self.folder_path = folder_path
        self.save_results = save_results
        self.yolo_detector = YOLO(os.path.join('../models/yolo_models', f'yolov8{yolo_model}-pose.onnx'))
        self.kpts_confidence = 0.5
        self.frame = None
        self.plot_data = PlotPoseData()

    async def plot_skeleton_kpts(self, kpts, limbs, pose_limb_color, pose_kpt_color):
        for p_id, point in enumerate(kpts):
            x_coord, y_coord, conf = point
            if conf < self.kpts_confidence:
                continue
            r, g, b = pose_kpt_color[p_id]
            cv2.circle(self.frame, (int(x_coord), int(y_coord)), 5, (int(r), int(g), int(b)), -1)
        for sk_id, sk in enumerate(limbs):
            r, g, b = pose_limb_color[sk_id]
            if kpts[sk[0]][2] < self.kpts_confidence or kpts[sk[1]][2] < self.kpts_confidence:
                continue
            pos1 = int(kpts[sk[0]][0]), int(kpts[sk[0]][1])
            pos2 = int(kpts[sk[1]][0]), int(kpts[sk[1]][1])
            cv2.line(self.frame, pos1, pos2, (int(r), int(g), int(b)), thickness=2)

    @staticmethod
    def get_frame_centroids(detections, conf: float, frame_shape: tuple):
        width, height = frame_shape
        xyxy = [d.boxes.xyxy.data.numpy()[0] for d in detections[0] if float(d.boxes.conf.data) > conf]
        centroids = [((x1 / width + x2 / width) / 2, (y1 / height + y2 / height) / 2) for x1, y1, x2, y2 in xyxy]
        return centroids

    @staticmethod
    def get_kpts_centroid(detection, kpts, conf: float, frame_shape: tuple):
        """
        Возвращает центроид по ключевым точкам, где:
        5 - левое плечо, 6 - правое плечо, 11 - левое бедро, 12 - правое бедро.
        Если видны все 4 точки, то вычисляется центроид относительно них;
        если видны только верхние две, то центроидом считается середина отрезка между плечами;
        если ни одни из точек не видна, то центроид находится, исходя из bbox'а
        """
        points = np.concatenate([[kpts[:, :-1][5]], [kpts[:, :-1][6]], [kpts[:, :-1][11]], [kpts[:, :-1][12]]])
        points = points / frame_shape  # в относительные координаты
        confs = np.array([kpts[:, -1][5], kpts[:, -1][6], kpts[:, -1][11], kpts[:, -1][12]])
        if all(confs > conf):
            return points.sum(axis=0) / len(points)
        elif any(confs[-2:] < conf) and all(confs[:-2] > conf):
            return points[:-2].sum(axis=0) / 2
        else:
            xyxy = np.array(
                [detection.boxes.xyxy.data.numpy()[0][:-2], detection.boxes.xyxy.data.numpy()[0][-2:]]
            ) / frame_shape
            return xyxy.sum(axis=0) / 2

    async def make_centroids(self):
        X = []
        for im_num, image_name in enumerate(
                tqdm(os.listdir(self.folder_path), desc="Processing files", unit="files", ncols=75, colour='#37B6BD')
        ):
            self.frame = cv2.imread(os.path.join(self.folder_path, image_name))
            detections = self.yolo_detector.predict(self.frame, classes=[0], stream=False,
                                                    conf=self.kpts_confidence, verbose=False)
            if len(detections[0]) == 0:
                continue
            frame_shape = self.frame.shape[:-1][::-1]
            tasks = []
            centroids = []
            for detection in detections[0]:
                kpts = detection.keypoints.data.numpy()[0]
                centroid = self.get_kpts_centroid(detection, kpts, self.kpts_confidence, frame_shape)
                if any(centroid >= 1):
                    continue
                centroids.append(centroid)
                cv2.circle(self.frame, tuple((centroid * frame_shape).astype(int)), 7, (0, 255, 0), 2)
                tasks.append(asyncio.create_task(
                    self.plot_skeleton_kpts(detection.keypoints.data.numpy()[0], self.plot_data.limbs,
                                            self.plot_data.pose_limb_color, self.plot_data.pose_kpt_color)
                ))
            for task in tasks:
                await task
            if centroids:
                X.append(centroids)
            if self.save_results:
                cv2.imwrite(os.path.join('../prepared_images', f'{im_num}.png'), self.frame)
        return np.concatenate(X)

    async def clusters_search(self):
        wcss = []
        # X = await self.make_centroids()
        # torch.save(X, 'centroids_kpts.pt')
        X = torch.load('../models/kmeans_train_data/centroids_kpts.pt')
        for i in range(1, 11):
            kmeans = KMeans(n_clusters=i, init='k-means++', max_iter=300, n_init=10, random_state=42)
            kmeans.fit(X)
            wcss.append(kmeans.inertia_)

        plt.figure(figsize=(10, 6))
        plt.plot(range(1, 11), wcss)
        plt.title('Выбор количества кластеров методом локтя')
        plt.xlabel('Количество кластеров')
        plt.ylabel('WCSS')
        plt.grid()
        plt.show()

    @staticmethod
    def kmeans_test():
        X = torch.load('../models/kmeans_train_data/centroids_kpts.pt')
        kmeans = KMeans(n_clusters=5, init='k-means++', max_iter=300, n_init=10, random_state=42)
        y_pred = kmeans.fit_predict(X)
        plt.figure(figsize=(10, 6))
        plt.scatter(X[:, 0], X[:, 1], c=y_pred, cmap='Paired')
        plt.scatter(kmeans.cluster_centers_[:, 0], kmeans.cluster_centers_[:, 1], s=150, c='red', marker='^',
                    label='Centroids')
        plt.legend(loc='upper right')
        plt.show()

if __name__ == '__main__':
    proc = DistProcessing('x', '../images')
    # asyncio.run(proc.clusters_search())
    proc.kmeans_test()


from operator import index
from bs4 import BeautifulSoup
import requests
import re
import pandas as pd
import os ,sys
import time
import random
import json
import tqdm
from datetime import datetime

sys.path.append('/home/aift-ml/workspace/lm/KoELECTRA/finetune/KoBigBird-master/finetune/ai_analyst_api/crawler')

hdr = {'User-Agent' :'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/101.0.4951.67 Safari/537.36'}

class Crawler:
    def crawler(self, company_name, company_code, days):
        # print("-----------------ZOOM 금융 신문기사 수집을 시작합니다. -----------------")
        page = 1
        total_source_result = list()
        total_date_result = list()
        total_title_result = list()
        total_link_result = list()
        total_content_result = list()
        days_list = []
        stop_flag = False
        while page <= int(100):
            # print(company_code)
            url = 'https://invest.zum.com/api/domestic/stock/'+ company_code +'/news?'+'page='+str(page)+'&size=10'
            # print(url)
            time.sleep(random.uniform(4,6))
            source_code = requests.get(url, headers=hdr).text
            html = BeautifulSoup(source_code, "lxml")
            data = html.select_one('p').text
            json_data = json.loads(data)
            url_list = []
            title_list = []
            date_list = []
            source_list = []

            for i in json_data['newses']:
                url = i['id'] # 뉴스  링크
                title = i['title'] # 뉴스 제목
                source = i['mediaName']  # 뉴스 매체
                date = i['registerDateTime'] # 뉴스 날짜
                date = date.replace('T',' ').replace('-','.')[:16]
                url_list.append(url)
                title_list.append(title)
                date_list.append(date)
                source_list.append(source)

            total_source_result.extend(source_list)
            total_date_result.extend(date_list)
            total_title_result.extend(title_list)
            total_link_result.extend(url_list)

            for idx, date in enumerate(total_date_result):
                date = datetime.strptime(str(date).strip(), '%Y.%m.%d %H:%M').date()
                if date not in days_list:
                    days_list.append(date)
                    first_date = days_list[0]
                    last_date = days_list[-1]
                    # print((first_date-last_date).days)

                    if int((first_date-last_date).days) > int(days): ### 날짜 이상 들어오면 스탑하고 그전 까지만
                        remove_index = idx
                        total_source_result = total_source_result[:remove_index]
                        total_date_result = total_date_result[:remove_index]
                        total_title_result = total_title_result[:remove_index]
                        total_link_result = total_link_result[:remove_index]
                        stop_flag = True

            # print(days_list)
            first_date = days_list[0]
            last_date = days_list[-1]
            date_tuple = (first_date, last_date)
            # print(len(days_list))
            if stop_flag == True:
                break
            page += 1

        # print("크롤링 할 링크 수: ", remove_index)
        content_list = []
        if len(total_link_result) ==0:
            # print("해당 기업의 뉴스기사가 없습니다.")
            pass
        else:
            for add in total_link_result:
                time.sleep(random.uniform(4,6))
                res = requests.get(add, headers=hdr)
                soup = BeautifulSoup(res.text, 'lxml')
                content = soup.select_one('div.article_body').text
                # print(content)
                if content:
                    content = content.strip()
                else:
                    content = ""
                content = content.replace("\n"," ")
                content = re.sub(' +', ' ', content)
                content_list.append(content)

            temp_result= {"날짜" : total_date_result, "언론사" : total_source_result, "기사제목" : total_title_result, "내용" : content_list, "링크" : total_link_result}
            temp_df = pd.DataFrame(temp_result)
            temp_df['회사명'] = company_name
            temp_df['종목코드'] = company_code
            # print("크롤링 기사 개수: ",len(temp_df))

                # content_list = []
                # for link in links:
                #     add = 'https://finance.naver.com' + link.find('a')['href']
                #     link_result.append(add)
                #     res = requests.get(add)
                #     time.sleep(4)
                #     soup = BeautifulSoup(res.text, 'lxml')6
                #     content = soup.select_one('#news_read')
                #     content = content.text.strip()
                #     content_list.append(content)
            df_result = temp_df
            # print("----------------- ZOOM 투자 신문기사 수집을 완료하였습니다. -----------------")
            # print('./result/0714/crwaling_' + str(company_name) + '.xlsx')
            # df_result.to_excel('./result/crwaling_' + str(company_name) + '.xlsx', encoding='utf-8')
        return df_result, date_tuple

        # # 변수들 합쳐서 해당 디렉토리에 csv파일로 저장하기
        # print(page)
        # result= {"날짜" : total_date_result, "언론사" : total_source_result, "기사제목" : total_title_result,"내용":total_content_result, "링크" : total_link_result}
        # df_result = pd.DataFrame(result)

        # print("다운 받고 있습니다------")
        # df_result.to_excel('./result/crwaling_' + str(company_name) + '.xlsx', encoding='utf-8')

    # 종목 리스트 파일 열기
    # 회사명을 종목코드로 변환

    def convert_to_code(self, company, days):
        data = pd.read_csv('/home/aift-ml/workspace/lm/KoELECTRA/finetune/KoBigBird-master/finetune/ai_analyst_api/crawler/resource/company_list.txt', dtype=str, sep='\t')   # 종목코드 추출
        company_name = data['회사명']
        keys = [i for i in company_name]    #데이터프레임에서 리스트로 바꾸기

        company_code = data['종목코드']
        values = [j for j in company_code]
        values_re = []
        for v in values:
            if len(str(v)) == 6:
                values_re.append(v)
            elif len(str(v)) == 5:
                v = '0'+ str(v)
                values_re.append(v)
            elif len(str(v)) == 4:
                v = '00'+ str(v)
                values_re.append(v)
            elif len(str(v)) == 3:
                v = '000'+ str(v)
                values_re.append(v)

        dict_result = dict(zip(keys, values_re))  # 딕셔너리 형태로 회사이름과 종목코드 묶기
        pattern = '[a-zA-Z가-힣]+'

        if bool(re.match(pattern, company)) == True:         # Input에 이름으로 넣었을 때
            company_code = dict_result.get(str(company))
            company_name = company
            # print(company_name,"     ",company_code)
            result = self.crawler(company_name,company_code, days)

        else:                                                # Input에 종목코드로 넣었을 때
            company_code = str(company)
            company_name = company
            # print(company_name,"     ",company_code)
            result = self.crawler(company_name ,company_code, days)

        return result


from django.conf.urls import patterns, include, url
from django.conf import settings
from django.conf.urls.static import static
# Uncomment the next two lines to enable the admin:
from django.contrib import admin
from main import views
admin.autodiscover()

urlpatterns = patterns('',
    # Examples:
    # url(r'^$', 'SportNewsApp.views.home', name='home'),
    # url(r'^SportNewsApp/', include('SportNewsApp.foo.urls')),

    url(r'^$', include('main.urls')),
    url(r'^news/', include('main.urls')),
    url(r'^accounts/', include('main.urls')),
    url(r'^info/', views.info, name='info'),

    # Uncomment the admin/doc line below to enable admin documentation:
    url(r'^admin/doc/', include('django.contrib.admindocs.urls')),

    # Uncomment the next line to enable the admin:
    url(r'^admin/', include(admin.site.urls)),

    url(r'^statistics/', include('statistics.urls')),

    url(r'^static/(?P<path>.*)$', 'django.views.static.serve', {'document_root': settings.STATIC_ROOT } ),
)

if settings.DEBUG:
    urlpatterns += patterns('',
        url(r'^media/(?P<path>.*)$', 'django.views.static.serve', {
        'document_root': settings.MEDIA_ROOT,
    }),
    )


#!/usr/bin/env python3

import os
import subprocess
import sys

import yaml

def read_config(path):
    with open(path, "r") as fp:
        return yaml.load(fp, Loader=yaml.FullLoader)

def run(scenario_py, cname, cvars, output_dir):
    cmd = [
        "python",
        scenario_py,
        # necessary to copy PYTHONPATH for venvs
        "--copy-env",
        "--append",
        os.path.join(output_dir, "results.json"),
        "--name",
        cname,
    ]
    for cvarname, cvarval in cvars.items():
        cmd.append("--{}".format(cvarname))
        cmd.append(str(cvarval))
    proc = subprocess.Popen(cmd)
    proc.wait()

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: {} <output dir>".format(sys.argv[0]))
        sys.exit(1)

    output_dir = sys.argv[1]
    print("Saving results to {}".format(output_dir))
    config = read_config("config.yaml")
    for cname, cvars in config.items():
        run("scenario.py", cname, cvars, output_dir)


"""
Download images related to a room
"""
import requests
import json
from typing import Dict
from pathlib import Path
from tqdm.auto import tqdm
import dask.multiprocessing
import dask.bag as db
import time

def download_photos(listing_id: int, download_dir: Path):
    path = Path(download_dir) / str(listing_id) / "photos"
    path.mkdir(exist_ok=True, parents=True)

    with open(path / "photos.json", "w") as fid:
        json.dump(photos, fid)

if __name__ == "__main__":

    with open("rooms.json") as fid:
        rooms = json.load(fid)

    download = Path("download")

    def is_not_done(listing_id):
        return not (download / str(listing_id) / "photos.json").is_file()

    errors = []

    for room in tqdm(rooms):
        if is_not_done(room):
            try:
                download_room(room, download)
            except:
                print(f"issue with room {room}")
                errors.append(room)
                with open("errors-room.json", "w") as fid:
                    json.dump(errors, fid)
    # dask.config.set(scheduler="processes", num_workers=20)
    # db.from_sequence(rooms).map(lambda l_id: download_room(l_id, download)).compute()


from datetime import date, datetime

def convert_to_date(date_input):
    """datetime形式、date形式、またはstr形式の日時情報を入力として、date形式を返す関数"""
    if date_input is None:
        print("警告: 入力がNoneです。1970-01-01を返します。")
        return date(1970, 1, 1)

    if isinstance(date_input, datetime):
        return date_input.date()
    elif isinstance(date_input, date):
        return date_input
    elif isinstance(date_input, str):
        try:
            parsed_date = datetime.strptime(date_input, "%Y-%m-%d")
        except ValueError:
            try:
                parsed_date = datetime.strptime(date_input, "%Y-%m-%d %H:%M:%S")
            except ValueError:
                print("無効な日付形式です。YYYY-MM-DDまたはYYYY-MM-DD hh:mm:ss形式で入力してください。")
                return None
        return parsed_date.date()
    else:
        print("無効な入力です。datetime形式、date形式、またはstr形式を使用してください。")
        return None

import unittest
from datetime import datetime, date

class TestConvertToDate(unittest.TestCase):
    def test_convert_to_date_with_datetime(self):
        input_date = datetime(2022, 1, 1, 12, 0, 0)
        expected_date = date(2022, 1, 1)
        self.assertEqual(convert_to_date(input_date), expected_date)

    def test_convert_to_date_with_date(self):
        input_date = date(2022, 1, 1)
        expected_date = date(2022, 1, 1)
        self.assertEqual(convert_to_date(input_date), expected_date)

    def test_convert_to_date_with_valid_string(self):
        input_date = "2022-01-01"
        expected_date = date(2022, 1, 1)
        self.assertEqual(convert_to_date(input_date), expected_date)

    def test_convert_to_date_with_valid_string_and_time(self):
        input_date = "2022-01-01 12:00:00"
        expected_date = date(2022, 1, 1)
        self.assertEqual(convert_to_date(input_date), expected_date)

    def test_convert_to_date_with_invalid_string(self):
        input_date = "invalid date"
        self.assertIsNone(convert_to_date(input_date))

    def test_convert_to_date_with_none(self):
        input_date = None
        expected_date = date(1970, 1, 1)
        self.assertEqual(convert_to_date(input_date), expected_date)

    def test_convert_to_date_with_invalid_type(self):
        input_date = 12345
        self.assertIsNone(convert_to_date(input_date))

if __name__ == "__main__":
    unittest.main()


from datetime import time

import datetime
import dill
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, FunctionTransformer
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import cross_val_score
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import MinMaxScaler
from sklearn.utils import resample
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC

def merge_data(sessions, hits):
    import pandas as pd
    sessions_df = sessions.copy()
    hits_df = hits.copy()
    df_hits_filtered = hits_df[((hits_df['event_action'] == 'sub_car_claim_click') | (
            hits_df['event_action'] == 'sub_car_claim_submit_click') | (
                                        hits_df['event_action'] == 'sub_open_dialog_click') | (
                                        hits_df['event_action'] == 'sub_custom_question_submit_click') | (
                                        hits_df['event_action'] == 'sub_call_number_click') | (
                                        hits_df['event_action'] == 'sub_callback_submit_click') | (
                                        hits_df['event_action'] == 'sub_submit_success') | (
                                        hits_df['event_action'] == 'sub_car_request_submit_click'))]
    hits_df = df_hits_filtered.groupby('session_id').count()
    merged_df = pd.merge(sessions_df, hits_df[['event_action']], on='session_id', how='left')
    merged_df.loc[merged_df['event_action'].isna(), 'event_action'] = 0
    merged_df.loc[merged_df['event_action'] > 0, 'event_action'] = 1
    return merged_df

def clean_df(sessions):
    def get_boundaries(data):
        q25 = data.quantile(0.25)
        q75 = data.quantile(0.75)
        iqr = q75 - q25
        boundaries = (round(q25 - 1.5 * iqr), round(q75 + 1.5 * iqr))
        return boundaries

    sessions_df = sessions.copy()
    space_mask = ((sessions_df['device_brand'] == '') | (sessions_df['device_brand'].str.isspace()))
    sessions_df.loc[space_mask, 'device_brand'] = None
    not_set_brand_mask = (sessions_df['device_brand'] == '(not set)')
    sessions_df.loc[not_set_brand_mask, 'device_brand'] = None
    not_set_os_mask = (sessions_df['device_os'] == '(not set)')
    sessions_df.loc[not_set_os_mask, 'device_os'] = None
    sessions_df.loc[not_set_os_mask, 'device_browser'] = None
    df_apple_sessions_mobile = ((sessions_df['device_brand'] == 'Apple') & (sessions_df['device_os'].isna()) & (
            sessions_df['device_category'] == 'mobile'))
    df_apple_sessions_tablet = ((sessions_df['device_brand'] == 'Apple') & (sessions_df['device_os'].isna()) & (
            sessions_df['device_category'] == 'tablet'))
    df_apple_sessions_desktop = ((sessions_df['device_brand'] == 'Apple') & (sessions_df['device_os'].isna()) & (
            sessions_df['device_category'] == 'desktop'))
    sessions_df.loc[df_apple_sessions_mobile, 'device_os'] = 'iOS'
    sessions_df.loc[df_apple_sessions_tablet, 'device_os'] = 'iOS'
    sessions_df.loc[df_apple_sessions_desktop, 'device_os'] = 'Macintosh'
    df_safari_sessions_mobile_os_na = (
            (sessions_df['device_browser'].str.contains('Safari')) & (sessions_df['device_os'].isna()) & (
            sessions_df['device_category'] == 'mobile'))
    df_safari_sessions_tablet_os_na = (
            (sessions_df['device_browser'].str.contains('Safari')) & (sessions_df['device_os'].isna()) & (
            sessions_df['device_category'] == 'tablet'))
    df_safari_sessions_desktop_os_na = (
            (sessions_df['device_browser'].str.contains('Safari')) & (sessions_df['device_os'].isna()) & (
            sessions_df['device_category'] == 'desktop'))
    sessions_df.loc[df_safari_sessions_mobile_os_na, 'device_os'] = 'iOS'
    sessions_df.loc[df_safari_sessions_tablet_os_na, 'device_os'] = 'iOS'
    sessions_df.loc[df_safari_sessions_desktop_os_na, 'device_os'] = 'Macintosh'
    df_safari_sessions_brand_na = (
            sessions_df['device_browser'].str.contains('Safari') & (sessions_df['device_brand'].isna()))
    apple_mac_mask = ((sessions_df['device_os'] == 'Macintosh') & (sessions_df['device_brand'].isna()))
    sessions_df.loc[df_safari_sessions_brand_na, 'device_brand'] = 'Apple'
    sessions_df.loc[apple_mac_mask, 'device_brand'] = 'Apple'
    sessions_df.loc[sessions_df['device_brand'].isna(), 'device_brand'] = 'Unknown'
    boundaries = get_boundaries(sessions_df['visit_number'])
    is_outlier = (sessions_df['visit_number'] < boundaries[0]) | (sessions_df['visit_number'] > boundaries[1])
    sessions_df.loc[is_outlier, 'visit_number'] = boundaries[1]
    boundaries = None
    sessions_df['visit_time'] = sessions_df['visit_time'].astype('str')
    sessions_df['event_action'] = sessions_df['event_action'].astype('int64')
    sessions_df = sessions_df[sessions_df['utm_source'].notna()]
    sessions_df = sessions_df[sessions_df['device_browser'].notna()]
    sessions_df = sessions_df[sessions_df['utm_campaign'].notna()]
    sessions_df = sessions_df[sessions_df['utm_adcontent'].notna()]
    return sessions_df

def decode_location(dataframe):
    from geopy.geocoders import Nominatim
    from geopy.extra.rate_limiter import RateLimiter
    import os.path
    geolocator = Nominatim(user_agent='my_application')
    geocode = RateLimiter(geolocator.geocode, min_delay_seconds=1)

    def encode_city_location(city_name):
        location = geocode(city_name, timeout=None)
        if not location:
            location = geocode('(not set)')
        return [location.latitude, location.longitude]

    def get_city_names_list(df_column):
        output = df_column.unique()
        return output

    def calculate_city_coords(cities_list):
        import json
        path = './model/data/city_coords.json'
        check_file = os.path.isfile(path)
        output_update = False
        if check_file:
            with open(path, 'r') as openfile:
                output = json.load(openfile)
        else:
            output_update = True
            output = dict()
        for city_name in cities_list:
            if city_name not in output:
                output[city_name] = encode_city_location(city_name)
        if output_update:
            with open(path, 'w') as outfile:
                json.dump(output, outfile)
        return output

    df_output = dataframe.copy()
    city_names_list = get_city_names_list(dataframe['geo_city'])
    city_coords_list = calculate_city_coords(city_names_list)
    for item in city_coords_list:
        df_output.loc[df_output['geo_city'] == item, ['city_lat', 'city_long']] = city_coords_list[item][0], \
            city_coords_list[item][1],
    df_output = df_output.drop([
        'geo_city',
    ], axis=1)
    return df_output

def decode_timestamps(dataframe):
    import pandas as pd

    def calculate_hour(income_string):
        splitted_string = income_string.split(':')
        hour = int(splitted_string[0])
        minutes = int(splitted_string[1])
        if minutes > 30:
            if hour == 23:
                hour = 0
            else:
                hour += 1
        return hour
    df_output = dataframe.copy()
    df_output['visit_date'] = pd.to_datetime(df_output['visit_date'])
    df_output['visit_dayofweek'] = df_output['visit_date'].dt.dayofweek
    df_output['visit_day'] = df_output['visit_date'].dt.day
    df_output['visit_month'] = df_output['visit_date'].dt.month
    df_output['visit_hour'] = df_output['visit_time'].apply(lambda x: calculate_hour(x))
    df_output = df_output.drop([
        'visit_date',
        'visit_time',
    ], axis=1)
    return df_output

def decode_device_sizes(dataframe):
    import pandas as pd

    def encode_pixels(income_row):
        income_string = income_row['device_screen_resolution']
        if income_string == '(not set)':
            return [10, 10]
        splitted_string = income_string.split('x')
        return [int(splitted_string[0]), int(splitted_string[1])]

    df_output = dataframe.copy()
    pixel_numbers = df_output.apply(lambda x: encode_pixels(x), axis=1, result_type='expand')
    pixel_numbers.columns = ['pixel_width', 'pixel_height']
    df_output = pd.concat([df_output.reset_index(drop=True), pixel_numbers.reset_index(drop=True)], axis=1)
    df_output = df_output.drop([
        'device_screen_resolution',
    ], axis=1)
    df_output['pixel_width'] = df_output['pixel_width'].apply(lambda x: x / 10)
    df_output['pixel_width'] = df_output['pixel_width'].round(0)
    df_output['pixel_height'] = df_output['pixel_height'].apply(lambda x: x / 10)
    df_output['pixel_height'] = df_output['pixel_height'].round(0)
    return df_output

def create_short_browser(dataframe):
    df_output = dataframe.copy()
    df_output['short_browser'] = df_output['device_browser'].apply(lambda x: x.split(' ')[0])
    df_output = df_output.drop([
        'device_browser',
    ], axis=1)
    return df_output

def downsampling_dataset(input_df):
    import pandas as pd
    initial_df = input_df.copy()
    df_majority = initial_df[initial_df['event_action'] == 0]
    df_minority = initial_df[initial_df['event_action'] == 1]
    df_minority_len = df_minority.shape[0]
    df_majority_downsampled = resample(df_majority,
                                       replace=False,
                                       n_samples=df_minority_len,
                                       random_state=852)
    df_downsampled = pd.concat([df_majority_downsampled, df_minority])
    return df_downsampled

def drop_unnecessary_columns(df):
    output = df.copy()
    output = output.drop([
        'device_model',
        'utm_keyword',
        'device_os',
        'session_id',
        'client_id',
        'geo_country',
    ], axis=1)
    return output

def main():
    print('Event Prediction Pipeline')
    df_sessions = pd.read_pickle('./model/data/ga_sessions.pkl')
    df_hits = pd.read_pickle('./model/data/ga_hits.pkl')
    merged_df = merge_data(df_sessions, df_hits)
    merged_df = downsampling_dataset(merged_df)
    merged_df = clean_df(merged_df)
    X = merged_df.drop(['event_action'], axis=1)
    y = merged_df['event_action']
    df_cleaner = Pipeline(steps=[
        ('drop_columns', FunctionTransformer(drop_unnecessary_columns)),
    ])

    feature_creation = Pipeline(steps=[
        ('create_short_browser', FunctionTransformer(create_short_browser)),
        ('decode_device_size', FunctionTransformer(decode_device_sizes)),
        ('decode_timestamp', FunctionTransformer(decode_timestamps)),
        ('decode_geo_location', FunctionTransformer(decode_location)),
    ])

    standart_scaling = Pipeline(steps=[
        ('std_scaler', StandardScaler()),
    ])

    min_max_scaling = Pipeline(steps=[
        ('mm_scaler', MinMaxScaler()),
    ])

    categorical_transformer = Pipeline(steps=[
        ('encoder', OneHotEncoder(handle_unknown='ignore')),
    ])

    column_transformer = ColumnTransformer(transformers=[
        (
            'numerical_std', standart_scaling,
            ['visit_number', 'visit_dayofweek', 'visit_day', 'visit_month', 'visit_hour',
             'pixel_width', 'pixel_height']),
        ('numerical_mm', min_max_scaling, ['city_lat', 'city_long']),
        ('categorical', categorical_transformer, [
            'utm_source',
            'utm_medium',
            'utm_campaign',
            'utm_adcontent',
            'device_category',
            'device_brand',
            'short_browser',
        ])
    ])

    preprocessor = Pipeline(steps=[
        ('filter', df_cleaner),
        ('feature_creator', feature_creation),
        ('column_transformer', column_transformer),
    ])

    models = (
        MLPClassifier(random_state=42, max_iter=19000, hidden_layer_sizes=5, learning_rate='adaptive', solver='lbfgs'),
        RandomForestClassifier(random_state=42, bootstrap=True, n_estimators=200, min_samples_split=15),
        SVC(),
    )

    best_score = .0
    best_pipe = None
    for model in models:
        pipe = Pipeline(steps=[
            ('preprocessor', preprocessor),
            ('classifier', model),
        ])
        score = cross_val_score(pipe, X, y, cv=4, scoring='roc_auc')
        print(f'model: {type(model).__name__}, roc_auc_mean: {score.mean():.2f}, roc_auc_std: {score.std():.2f}')

        if score.mean() > best_score:
            best_score = score.mean()
            best_pipe = pipe

    best_pipe.fit(X, y)
    print(f'best model: {type(best_pipe.named_steps["classifier"]).__name__}, roc_auc: {best_score:.2f}')
    with open('model/best_pipe.pkl', 'wb') as file:
        dill.dump({
            'model': best_pipe,
            'metadata': {
                'name': 'Sberauto event prediction pipeline',
                'author': 'Roman Suvorkov',
                'version': 1,
                'date': datetime.datetime.now(),
                'type': type(best_pipe.named_steps['classifier']).__name__,
                'roc_auc': best_score,
            }
        }, file)

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    main()


import requests
from bs4 import BeautifulSoup

class searchLinkedIn:
	def __init__(self):
		server = "encrypted.google.com"
		limit = "100"

		self.linkedin_list = ["| LinkedIn", "on LinkedIn", "- LinkedIn", "LinkedIn"]
		self.google_search = "https://"+server+"/search?num="+limit+"&q=%s %s insite:fr.linkedin.com/in"

	def search(self, searching, city):
		google_search = self.google_search % (searching, city)
		linkedin_list = self.linkedin_list

		employee_list = []
		profile_list = []

		req = requests.get(google_search)
		status_code = req.status_code

		if status_code == 200:
			html = BeautifulSoup(req.text, "html.parser")
			results = html.find_all('div', { 'class' : 'g' })

			for res in results:
				employee = res.find('h3', { 'class' : 'r' }).text

				if "LinkedIn" in employee:
					for l in linkedin_list:
						if l in employee:
							employee = employee.replace(l, "")
							employee_list.append(employee)

					profile = res.find('cite').text
					profileSplit = profile.split("/")
					profile = "/in/"+profileSplit[len(profileSplit) -1]
					profile_list.append(profile)

				if len(employee_list) > 0:
					found = len(employee_list)
					employees = employee_list
					profiles = profile_list

		self.found = len(employee_list)
		self.employees = employee_list
		self.profiles = profile_list


"""==================================================================================================="""
### Basic Libraries
import warnings
warnings.filterwarnings("ignore")
import os, sys, numpy as np, argparse, imp, datetime, pandas as pd, copy
import time, pickle as pkl, random, json, collections

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
from tqdm import tqdm

### Repository-specific Libraries
import parameters    as par
import utilities.misc as misc

"""==================================================================================================="""
parser = argparse.ArgumentParser()

parser = par.basic_training_parameters(parser)
parser = par.batch_creation_parameters(parser)
parser = par.batchmining_specific_parameters(parser)
parser = par.loss_specific_parameters(parser)
parser = par.wandb_parameters(parser)
parser = par.s2sd_parameters(parser)
parser = par.maxentropy_parameters(parser)
parser.add_argument('--data_hardness', default=1, type=int, help='OOD-ness of datasets. Ranges/Default id are [CUB200-2011, (1-9), 3], [CARS196, (1-8), 1], [SOP, (1-8), 1]')
parser.add_argument('--new_test_dataset', default='None', type=str, help='Optionally evaluate on the test dataset of a different dataset.')
##### Read in parameters
opt = parser.parse_args()

"""==================================================================================================="""
### The following setting is useful when logging to wandb and running multiple seeds per setup:
### By setting the savename to <group_plus_seed>, the savename will instead comprise the group and the seed!
if opt.savename=='group_plus_seed':
    if opt.log_online:
        opt.savename = opt.group+'_s{}'.format(opt.seed)
    else:
        opt.savename = ''

### Easy checkpointing
start_from_checkpoint = False
if opt.checkpoint:
    init_save_path = copy.deepcopy(opt.save_path)
    checkpath      = opt.save_path+'/'+opt.dataset+'/'+opt.savename
    if os.path.exists(checkpath):
        if any(['checkpoint.pth.tar'==x for x in os.listdir(checkpath)]):
            load_opt = pkl.load(open(checkpath+'/hypa.pkl', 'rb'))
            new_run  = False
            for key,item in vars(opt).items():
                if key!='load_to_ram':
                    if item!=vars(load_opt)[key] and key not in ['source_path', 'save_path', 'completed']:
                        new_run = True

            if not new_run:
                opt             = load_opt
                opt.save_path   = init_save_path
                opt.source_path = '/'.join(opt.source_path.split('/')[:-1])
                start_from_checkpoint = True

if opt.completed:
    print('\n\nTraining Run already completed!')
    exit()

### If wandb-logging is turned on, initialize the wandb-run here:
if opt.log_online:
    import wandb
    _ = os.system('wandb login {}'.format(opt.wandb_key))
    os.environ['WANDB_API_KEY'] = opt.wandb_key
    if not start_from_checkpoint:
        opt.unique_run_id = wandb.util.generate_id()
    wandb.init(id=opt.unique_run_id, resume='allow', project=opt.project, group=opt.group, name=opt.savename, dir=opt.source_path)
    wandb.config.update(opt)

"""==================================================================================================="""
### Load Remaining Libraries that need to be loaded after wandb
import torch, torch.nn as nn, torch.nn.functional as F
import torch.multiprocessing
import torchvision
torch.multiprocessing.set_sharing_strategy('file_system')
import architectures as archs
import datasampler   as dsamplers
import datasets      as dataset_library
import criteria      as criteria
import metrics       as metrics
import batchminer    as bmine
import evaluation    as eval
from utilities import misc
from utilities import logger

"""==================================================================================================="""
full_training_start_time = time.time()

opt.source_path += '/'+opt.dataset
opt.save_path   += '/'+opt.dataset

#Assert that the construction of the batch makes sense, i.e. the division into class-subclusters.
assert not opt.bs%opt.samples_per_class, 'Batchsize needs to fit number of samples per class for distance sampling and margin/triplet loss!'

opt.pretrained      = not opt.not_pretrained
opt.evaluate_on_gpu = not opt.evaluate_on_cpu

os.environ["CUDA_DEVICE_ORDER"]   ="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]= str(opt.gpu[0])

if start_from_checkpoint:
    chkp_data = torch.load(opt.save_path+'/'+opt.savename+'/checkpoint.pth.tar')
misc.set_seed(opt.seed)

"""==================================================================================================="""
##################### NETWORK SETUP ##################
opt.device = torch.device('cuda')
model      = archs.select(opt.arch, opt)

if opt.fc_lr<0:
    to_optim   = [{'params':model.parameters(),'lr':opt.lr,'weight_decay':opt.decay}]
else:
    all_but_fc_params = [x[-1] for x in list(filter(lambda x: 'last_linear' not in x[0], model.named_parameters()))]
    fc_params         = model.model.last_linear.parameters()
    to_optim          = [{'params':all_but_fc_params,'lr':opt.lr,'weight_decay':opt.decay},
                         {'params':fc_params,'lr':opt.fc_lr,'weight_decay':opt.decay}]

if start_from_checkpoint:
    opt       = chkp_data['opt']
    _ = model.load_state_dict(chkp_data['state_dict'])

_  = model.to(opt.device)

"""============================================================================"""
# Code needed for Uniform-Prior regularization (Sinha et al., 2020)
class Discriminator(nn.Module):
    def __init__(self, in_dim, latent):
        super(Discriminator, self).__init__()
        self.in_dim = in_dim
        self.fc1 = nn.Linear(in_dim, latent)
        self.fc2 = nn.Linear(latent, latent)
        self.fc3 = nn.Linear(latent, 1)
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = torch.sigmoid(self.fc3(x))
        return x

if opt.with_entropy:
    discr = Discriminator(opt.maxentropy_chunksize, opt.maxentropy_latent)
    if start_from_checkpoint:
        _ = discr.load_state_dict(chkp_data['aux']['discr_state_dict'])
    _ = discr.to(opt.device)

    optimizer_discr     = torch.optim.Adam(discr.parameters(), lr=opt.maxentropy_lrmulti*opt.lr, betas=(0.5, 0.999))
    if start_from_checkpoint:
        optimizer_discr.load_state_dict(chkp_data['aux']['optimizer_discr_state_dict'])
    scheduler_discr = torch.optim.lr_scheduler.MultiStepLR(optimizer_discr, milestones=opt.maxentropy_tau, gamma=opt.maxentropy_gamma)
    if start_from_checkpoint:
        scheduler_discr.load_state_dict(chkp_data['aux']['scheduler_discr_state_dict'])

"""============================================================================"""
#################### DATALOADER SETUPS ##################
datasets    = dataset_library.select(opt.dataset, opt, opt.source_path)
if hasattr(model, 'specific_normalization'):
    for dataset_type in datasets.keys():
        if datasets[dataset_type] is not None:
            datasets[dataset_type].normal_transform_list[-1] = model.specific_normalization
            datasets[dataset_type].normal_transform = torchvision.transforms.Compose(
                datasets[dataset_type].normal_transform_list
            )

if opt.new_test_dataset != 'None':
    init_dataset = copy.deepcopy(opt.dataset)
    init_sourcepath = copy.deepcopy(opt.source_path)
    opt.source_path = opt.source_path.replace('/' + opt.dataset, '/' + opt.new_test_dataset)
    opt.dataset = opt.new_test_dataset
    datasets_temp = dataset_library.select(opt.new_test_dataset, opt, opt.source_path)
    datasets['testing'] = datasets_temp['testing']
    opt.dataset = init_dataset
    opt.source_path = init_sourcepath

dataloaders = {}
dataloaders['evaluation'] = torch.utils.data.DataLoader(datasets['evaluation'], num_workers=opt.kernels, batch_size=opt.bs, shuffle=False)
dataloaders['testing']    = torch.utils.data.DataLoader(datasets['testing'],    num_workers=opt.kernels, batch_size=opt.bs, shuffle=False)

if opt.use_tv_split:
    dataloaders['validation'] = torch.utils.data.DataLoader(datasets['validation'], num_workers=opt.kernels, batch_size=opt.bs,shuffle=False)

train_data_sampler      = dsamplers.select(opt.data_sampler, opt, datasets['training'].image_dict, datasets['training'].image_list)
if train_data_sampler.requires_storage:
    train_data_sampler.create_storage(dataloaders['evaluation'], model, opt.device)

dataloaders['training'] = torch.utils.data.DataLoader(datasets['training'], num_workers=opt.kernels, batch_sampler=train_data_sampler)

opt.n_classes  = len(dataloaders['training'].dataset.avail_classes)
opt.n_test_classes = len(dataloaders['testing'].dataset.avail_classes)

"""============================================================================"""
#################### CREATE LOGGING FILES ###############
sub_loggers = ['Train', 'Test', 'Model Grad']
if opt.use_tv_split: sub_loggers.append('Val')

if start_from_checkpoint:
    LOG = chkp_data['progress']
else:
    LOG = logger.LOGGER(opt, sub_loggers=sub_loggers, start_new=True, log_online=opt.log_online)

"""============================================================================"""
#################### Basic DML Loss Setup ####################
batchminer   = bmine.select(opt.batch_mining, opt)
criterion, to_optim = criteria.select(opt.loss, opt, to_optim, batchminer)

if start_from_checkpoint:
    _ = criterion.load_state_dict(chkp_data['aux']['criterion'][-1])

_ = criterion.to(opt.device)

if 'criterion' in train_data_sampler.name:
    train_data_sampler.internal_criterion = criterion

"""============================================================================"""
#################### OPTIM SETUP ####################
if opt.optim == 'adam':
    optimizer    = torch.optim.Adam(to_optim)
else:
    raise Exception('Optimizer <{}> not available!'.format(opt.optim))

if start_from_checkpoint:
    _ = optimizer.load_state_dict(chkp_data['aux']['optimizer'])

scheduler    = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=opt.tau, gamma=opt.gamma)
if start_from_checkpoint:
    scheduler.load_state_dict(chkp_data['aux']['scheduler'])

"""============================================================================"""
#################### METRIC COMPUTER ####################
metric_computer = metrics.MetricComputer(opt.evaluation_metrics, opt)

"""============================================================================"""
################### Summary #########################3
chkp_text  = 'RESTARTED FROM CHKP: {} {}\n'.format(start_from_checkpoint, '- at epoch {}'.format(opt.epoch+1) if 'epoch' in vars(opt) else '')
data_text  = 'Dataset:\t {}'.format(opt.dataset.upper())
setup_text = 'Objective:\t {}'.format(opt.loss.upper())
miner_text = 'Batchminer:\t {}'.format(opt.batch_mining if criterion.REQUIRES_BATCHMINER else 'N/A')
arch_text  = 'Backbone:\t {} (#weights: {})'.format(opt.arch.upper(), misc.gimme_params(model))
summary    = chkp_text+'\n'+data_text+'\n'+setup_text+'\n'+miner_text+'\n'+arch_text
print(summary)

"""============================================================================"""
################### SCRIPT MAIN ##########################
print('\n-----\n')

iter_count = 0
include_aux_augmentations = dataloaders['training'].dataset.include_aux_augmentations
loss_args  = {'batch':None, 'labels':None, 'batch_features':None, 'f_embed':None}
init_lrs = [param_group['lr'] for param_group in optimizer.param_groups[1:]]

if not start_from_checkpoint:
    opt.epoch  = 0
else:
    if 'epoch' in vars(opt):
        opt.epoch += 1
epochs = range(opt.epoch, opt.n_epochs)

for epoch in epochs:
    opt.epoch = epoch
    epoch_start_time = time.time()
    #Set seeds for each epoch - this ensures reproducibility after resumption.
    misc.set_seed(opt.n_epochs*opt.seed+epoch)

    ### Scheduling Changes specifically for cosine scheduling
    if opt.scheduler!='none':
        print('Running with learning rates {}...'.format(' | '.join('{}'.format(x['lr']) for x in optimizer.param_groups)))

    ### Train one epoch
    start = time.time()
    _ = model.train()

    loss_collect = []
    data_iterator = tqdm(dataloaders['training'], desc='Epoch {} Training...'.format(epoch))

    for i,out in enumerate(data_iterator):
        class_labels, input, input_indices = out

        input      = input.to(opt.device)
        model_args = {'x':input.to(opt.device), 'warmup':epoch<opt.warmup}

        out_dict  = model(**model_args)
        embeds, avg_features, features, extra_embeds = [out_dict[key] for key in ['embeds', 'avg_features', 'features', 'extra_embeds']]

        loss_args['input_batch']    = input
        loss_args['batch']          = embeds
        loss_args['extra_batches']  = extra_embeds
        loss_args['labels']         = class_labels
        loss_args['f_embed']        = model.model.last_linear
        loss_args['batch_features'] = features
        loss_args['avg_batch_features'] = avg_features
        loss_args['model']          = model
        loss_args['optim']          = optimizer

        loss      = criterion(**loss_args)

        if opt.with_entropy:
            # Apply adversarial uniformity regularization.
            if iter_count % opt.maxentropy_iter == 0:
                # DSC STEP
                embeddings = avg_features.view(-1, opt.maxentropy_chunksize).detach()
                z          = torch.rand(embeddings.size(0), opt.maxentropy_chunksize).to(opt.device) * 2. - 1.
                ones       = torch.ones(embeddings.size(0), 1).to(opt.device)
                preds_real = discr(z)
                real_loss  = F.binary_cross_entropy(preds_real, ones)

                zeros      = torch.zeros(embeddings.size(0), 1).to(opt.device)
                preds_fake = discr(embeddings.detach())
                fake_loss  = F.binary_cross_entropy(preds_fake, zeros)
                dsc_loss   = real_loss + fake_loss

                optimizer_discr.zero_grad()
                dsc_loss.backward()
                optimizer_discr.step()

            embeddings = avg_features.view(-1, opt.maxentropy_chunksize)
            ones       = torch.ones(embeddings.size(0), 1).to(opt.device)
            preds_fake = discr(embeddings)
            adv_loss   = F.binary_cross_entropy(preds_fake, ones)
            loss += opt.maxentropy_w * adv_loss

        """========================="""
        optimizer.zero_grad()
        loss.backward()

        ### Compute Model Gradients and log them!
        grads = np.concatenate([p.grad.detach().cpu().numpy().flatten() for p in model.parameters() if p.grad is not None])
        grad_l2, grad_max  = np.mean(np.sqrt(np.mean(np.square(grads)))), np.mean(np.max(np.abs(grads)))

        LOG.progress_saver['Model Grad'].log('Grad L2',  grad_l2,  group='L2')
        LOG.progress_saver['Model Grad'].log('Grad Max', grad_max, group='Max')

        ### Update network weights!
        optimizer.step()

        ###
        loss_collect.append(loss.item())

        ###
        iter_count += 1

        data_iterator.set_postfix_str('Loss: {0:.4f}'.format(np.mean(loss_collect)))
        if i==len(dataloaders['training'])-1: data_iterator.set_description('Epoch (Train) {0}: Mean Loss [{1:.4f}]'.format(epoch, np.mean(loss_collect)))

    ####
    result_metrics = {'loss': np.mean(loss_collect)}

    ####
    LOG.progress_saver['Train'].log('epochs', epoch)
    for metricname, metricval in result_metrics.items():
        LOG.progress_saver['Train'].log(metricname, metricval)
    LOG.progress_saver['Train'].log('time', np.round(time.time()-start, 4))

    ### Learning Rate Scheduling Step
    if opt.scheduler != 'none':
        scheduler.step()
    if opt.with_entropy:
        scheduler_discr.step()

    """======================================="""
    ### Evaluate Metric for Training & Test (& Validation)
    _ = model.eval()
    print('\nComputing Testing Metrics...')

    aux_store = {'optimizer':optimizer.state_dict(), 'criterion':(criterion, criterion.state_dict()),
                 'scheduler':scheduler.state_dict(), 'dataloaders':dataloaders,
                 'discr_state_dict':discr.state_dict() if opt.checkpoint and opt.with_entropy else None,
                 'optimizer_discr_state_dict':optimizer_discr.state_dict() if opt.checkpoint and opt.with_entropy else None,
                 'scheduler_discr_state_dict':scheduler_discr.state_dict() if opt.checkpoint and opt.with_entropy else None,
                 'datasets':datasets, 'train_data_sampler':train_data_sampler}

    eval.evaluate(opt.dataset, LOG, metric_computer, [dataloaders['testing']], model, opt, opt.evaltypes, opt.device, log_key='Test', aux_store=aux_store)
    if opt.use_tv_split:
        print('\nComputing Validation Metrics...')
        eval.evaluate(opt.dataset, LOG, metric_computer, [dataloaders['validation']], model, opt, opt.evaltypes, opt.device, log_key='Val', aux_store=aux_store)

    if not opt.no_train_metrics:
        print('\nComputing Training Metrics...')
        eval.evaluate(opt.dataset, LOG, metric_computer, [dataloaders['evaluation']], model, opt, opt.evaltypes, opt.device, log_key='Train', aux_store=aux_store)

    import wandb
    LOG.update(all=True)

    eval.set_checkpoint(model, opt, LOG, LOG.prop.save_path+'/checkpoint.pth.tar', aux=aux_store)

    print('\nTotal Epoch Runtime: {0:4.2f}s'.format(time.time()-epoch_start_time))
    print('\n-----\n')

"""======================================================="""
opt.completed = True
pkl.dump(opt,open(opt.save_path+"/hypa.pkl","wb"))


#!usr/bin/env/python
# -*- coding: utf-8 -*-

import numpy as np
import cv2
import datetime

videoDir = "./"
cap = cv2.VideoCapture(0)
fourcc = cv2.cv.FOURCC(*'XVID')
now = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
videoName = videoDir + 'VID' + now + '.avi'
outWriter = cv2.VideoWriter(videoName,fourcc, 25.0, (640,480))
while(cap.isOpened()):
    ret, frame = cap.read()
    if ret==True:
        outWriter.write(frame)
        cv2.imshow('frame',frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
    	break

outWriter.release()
cv2.destroyAllWindows()


"""Some code from my Jupyter notebook that I wanted to run with PyCharm debugging!"""

import numpy as np
import matplotlib.pyplot as plt
import ocelot
import hdbscan
import pickle
import time

from pathlib import Path

# INITIAL SETUP
# Change this one
data_location = Path("/media/emily/Emilys External Disk1/data")

if not data_location.exists():
    raise ValueError("data location is invalid!")

# These should stay the same
location_mwsc = data_location / Path('sc_cats/mwsc_ii_catalogue')
location_dr2_10_ocs = data_location / Path('gaia_dr2_10ocs')
location_simulated_populations = data_location / Path('simulated_ocs/191122_logt_6_to_10_parsec_1.2s')

# A bit of automated handling to make a sorted list of locations
location_dr2_10_ocs_iterable = np.sort(list(location_dr2_10_ocs.glob('*.pickle')))

def read_oc(number: int):
    with open(location_dr2_10_ocs_iterable[number], 'rb') as handle:
        return pickle.load(handle)

# RUNNING CODE
# Disgusting code to reverse engineer the names
name_10_ocs = ["_".join(location_dr2_10_ocs_iterable[x].stem.rsplit("-")[0].rsplit("_")[1:]) for x in range(10)]

# Cycle over everyone!
start = 0
end = 5
output_dir_name = "500_stars_locked"

for i, a_name in enumerate(name_10_ocs[start:end], start=start):

    print(f"Working on {a_name}")

    # Get the data
    data_gaia = read_oc(i)
    data_gaia = ocelot.cluster.cut_dataset(data_gaia, parameter_cuts={'phot_g_mean_mag': [-np.inf, 18]})
    data_gaia = ocelot.cluster.recenter_dataset(
        data_gaia, center=(data_gaia['ra'].mean(), data_gaia['dec'].mean()))

    if np.any(data_gaia['ra'] > 270) and np.any(data_gaia['ra'] < 90):
        data_gaia['ra'] = np.where(data_gaia['ra'] < 180, data_gaia['ra'], data_gaia['ra'] - 360)

    # Rescale it
    data_rescaled = ocelot.cluster.rescale_dataset(
        data_gaia, columns_to_rescale=('lon', 'lat', 'pmlon', 'pmlat', 'parallax'))

    # Get labels and probs
    print("  doing clustering analysis")
    start_time = time.time()
    clusterer = hdbscan.HDBSCAN(min_cluster_size=10, min_samples=10, cluster_selection_method='leaf')
    labels = clusterer.fit_predict(data_rescaled)
    probabilities = clusterer.probabilities_
    persistences = clusterer.cluster_persistence_
    cluster_time = time.time() - start_time

    del clusterer

    # Grab all of our cluster significances for every cluster
    print("  doing significance analysis")

    start_time = time.time()
    significances, log_likelihoods = ocelot.verify.cluster_significance_test(
        data_rescaled,
        labels,
        min_samples=10,
        make_diagnostic_plots=True,
        plot_output_dir=Path(f"./{output_dir_name}/{a_name}"),
        knn_kwargs={'overcalculation_factor': 5, 'cluster_nn_distance_type': 'internal', 'n_jobs': 1,
                    'min_field_stars': 500, 'max_field_stars': 500},
        test_type='all',
    )
    sig_time = time.time() - start_time

    print(f"  cluster: {cluster_time:.2f}s -- sig: {sig_time:.2f}s")

    # Also want unique labels without -1
    unique_labels, counts = np.unique(labels, return_counts=True)
    good_labels = unique_labels != -1
    unique_labels = unique_labels[good_labels]
    counts = counts[good_labels]

    for a_cluster in unique_labels:
        print(f"  plotting {a_cluster+1} of {len(unique_labels)}")

        # Get significance and log likelihood in a good format
        sig = (
            f"lr: {significances['likelihood'][a_cluster]:.2f}  /  "
            f"t: {significances['welch_t'][a_cluster]:.2f}  /  "
            f"mw: {significances['mann_w'][a_cluster]:.2f} \n"
            f"k1+: {significances['ks_one+'][a_cluster]:.2f}  /  "
            f"k1-: {significances['ks_one-'][a_cluster]:.2f}  /  "
            f"diff: {significances['ks_one+'][a_cluster] - significances['ks_one-'][a_cluster]:.2f}"
            f"\nk2+: {significances['ks_two+'][a_cluster]:.2f}  /  "
            f"k2-: {significances['ks_two-'][a_cluster]:.2f}  /  "
            f"diff: {significances['ks_two+'][a_cluster] - significances['ks_two-'][a_cluster]:.2f}"
        )

        fig, ax = ocelot.plot.clustering_result(
            data_gaia, labels, [-2, -2, -2] + [a_cluster], probabilities, make_parallax_plot=True,
            cmd_plot_y_limits=[8, 18], cluster_marker_radius=(3, 3, 3, 3), dpi=300, show_figure=False,
            figure_size=(6, 6),
            figure_title=f"Field around {a_name}, cluster {a_cluster}\n"
                         f"  detected by HDBSCAN w/ m_clSize=40, m_Pts=10\n"
                         f"  significance: {sig}\n"
                         f"  persistence: {persistences[a_cluster]:.4f}\n"
                         f"  n_stars: {counts[a_cluster]}"
        )
        ax[0, 0].set_title(None)
        ax[0, 1].set_title(None)
        ax[1, 0].set_title(None)
        ax[1, 1].set_title(None)

        fig.subplots_adjust(wspace=0.3)

        fig.savefig(f"./{output_dir_name}/{a_name}/{a_cluster}_ocelot.png", bbox_inches="tight")

        plt.close(fig)


import argparse
import torch
import os
import random
import numpy as np
import itertools
import datetime
import gym
import json
import mujoco_py
from functools import reduce
from argparse import Namespace
from copy import deepcopy
from tensorboardX import SummaryWriter

import exp_utils
from gym_mm.envs.freerun import FreeRun
from algo.buffer import ReplayMemory
from algo.disc import PredTrainer
from algo.utils import convert_to_onehot, wrapped_obs
from config import parser_main

start_time, exp_seed = exp_utils.get_start_time()

# load args
customed_args = vars(parser.parse_args())

with open("mujoco_configs/train/{}.json".format(customed_args["scenario"])) as f:
    json_config = json.load(f)

args_dict = {}
args_dict.update(json_config["common"])
args_dict.update(json_config[customed_args["algo"]])
args_dict.update(json_config["dads_indie"])
args_dict.update(customed_args)
args = Namespace(**args_dict)
args.seed = exp_seed
print(args)
print("now is running: {}!".format(args.run))

# args = parser.parse_args()
np.random.seed(args.seed)
torch.manual_seed(args.seed)

if args.scenario == "AntCustom-v0":
    from gym_mm.envs.ant_custom_env import AntCustomEnv
    env = AntCustomEnv()
    reduced_obs = True
else:
    env = gym.make(args.scenario)
    reduced_obs = False
env.seed(args.seed)
obs_shape_list = env.observation_space.shape
obs_shape = reduce((lambda x,y: x*y), obs_shape_list)
action_space = env.action_space
discrete_action = hasattr(action_space, 'n')

from algo.sac import SACTrainer
trainers = [SACTrainer(obs_shape, action_space, args) for _ in range(args.num_modes)]
update_interval = 1
updates_per_step = 1
memories = [ReplayMemory(args.buffer_limit) for _ in range(args.num_modes)]
cache = ReplayMemory(args.buffer_limit)

# TensorboardX
logdir = 'logs/{}/dads_indie/dads_{}_{}_run{}_{}'.format(args.scenario, args.algo, args.scenario, args.run, datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
writer = SummaryWriter(logdir=logdir)
save_args(json_config, args, logdir)

if reduced_obs:
    dtrainer = PredTrainer(2, args)
else:
    dtrainer = PredTrainer(obs_shape, args)

src = 1
rc = 0
running_reward = 0
running_sr = 0
avg_length = 0
timestep = 0
updates = 0
scale = args.reward_scale
max_mean_sr = 0
input_amp = 100
for i_episode in itertools.count(1):
    obs = env.reset()
    episode_reward = 0
    episode_sr = 0
    episode_steps = 0
    label = np.random.randint(0, high=args.num_modes)
    l = np.array([label])
    for t in range(args.max_episode_len):
        if args.start_steps < timestep:
            action, logprob = trainers[label].act(obs)
            if len(action.shape) > 1:
                action = action[0]
            if discrete_action:
                action = action[0]
        else:
            action = env.action_space.sample()
            logprob = np.array([1.0])

        if timestep > args.start_steps:
            if timestep % update_interval == 0:
                for _ in range(args.updates_per_step):
                    label_batch, state_batch, action_batch, logprob_batch, reward_batch, next_state_batch, mask_batch = memories[label].sample(batch_size=args.batch_size)

                    c1_loss, c2_loss, p_loss, ent_loss, alpha = trainers[label].update_parameters((state_batch, action_batch, logprob_batch, reward_batch, next_state_batch, mask_batch), updates)
                    writer.add_scalar('loss/critic_1', c1_loss, updates)
                    writer.add_scalar('loss/critic_2', c2_loss, updates)
                    writer.add_scalar('loss/policy', p_loss, updates)
                    writer.add_scalar('loss/entropy_loss', ent_loss, updates)
                    writer.add_scalar('entropy_temprature/alpha', alpha, updates)

                    updates += 1
            label_batch, state_batch, next_state_batch = cache.sample(batch_size=args.disc_batch_size)
            state_delta_batch = next_state_batch - state_batch
            label_onehot_batch = convert_to_onehot(label_batch, args.num_modes)
            if reduced_obs:
                d_loss = dtrainer.update_parameters((state_batch[:,:2], label_onehot_batch, state_delta_batch[:,:2] * input_amp))
            else:
                d_loss = dtrainer.update_parameters((state_batch, label_onehot_batch, state_delta_batch * input_amp))
            writer.add_scalar('loss/disc', d_loss, timestep)

        new_obs, reward, done, _ = env.step(action.tolist())

        if timestep > args.start_steps:
            L = args.num_modes
            alt_labels = np.concatenate([np.arange(0, label), np.arange(label+1, L)])
            obs_delta = new_obs - obs
            if reduced_obs:
                logp = dtrainer.score(obs[:2], convert_to_onehot(l, args.num_modes), obs_delta[:2] * input_amp)
            else:
                logp = dtrainer.score(obs, convert_to_onehot(l, args.num_modes), obs_delta * input_amp)
            alt_obs = np.tile(obs, [L-1,1])
            alt_new_obs = np.tile(new_obs, [L-1,1])
            alt_obs_delta = alt_new_obs - alt_obs
            if reduced_obs:
                alt_logp = dtrainer.score(alt_obs[:,:2], convert_to_onehot(alt_labels, args.num_modes), alt_obs_delta[:,:2] * input_amp)
            else:
                alt_logp = dtrainer.score(alt_obs, convert_to_onehot(alt_labels, args.num_modes), alt_obs_delta * input_amp)

            writer.add_scalar('logp/logp', logp, timestep)
            writer.add_scalar('logp/alt_logp', alt_logp.mean(), timestep)
            writer.add_scalar('logp/alt_logp_max', alt_logp.max(), timestep)
            writer.add_scalar('bn/var', dtrainer.pred.output_bn.running_var.mean().item(), timestep)
            sr = np.log(L) - np.log(1+np.exp(np.clip(alt_logp - logp, -20, 1)).sum(axis=0))
        else:
            sr = 0.
        sr *= scale

        episode_reward += reward
        episode_sr += sr
        timestep += 1
        episode_steps += 1
        if hasattr(env, 'max_steps'):
            mask = 1 if episode_steps == env.max_steps else float(not done)
        else:
            mask = float(not done)

        memories[label].push((label, obs, action, logprob, sr * src + reward * rc, new_obs, mask))
        cache.push((l, obs, new_obs))
        obs = new_obs
        if done:
            break

    avg_length += (t+1)
    running_reward += episode_reward
    running_sr += episode_sr
    # env.render(message=str(label))
    writer.add_scalar('stats/episode_reward', episode_reward, i_episode)
    writer.add_scalar('stats/episode_sr', episode_sr, i_episode)
    writer.add_scalar('stats/episode_length', t, i_episode)

    if i_episode % args.log_interval == 0:
        print("Episode: {}, length: {}, reward: {}, sr: {}".format(i_episode, int(avg_length/args.log_interval),
            int(running_reward/args.log_interval), int(running_sr/args.log_interval)))
        if running_sr/args.log_interval > max_mean_sr:
            max_mean_sr = running_sr/args.log_interval
            for x in range(len(trainers)):
                trainers[x].save_model(args.scenario, prefix="models/{}/dads_indie/run{}/".format(args.scenario, args.run), suffix="dads_indie_{}".format(x), silent=True)
            print("Models saved to "+"models/{}/dads_indie/run{}".format(args.scenario, args.run))

        avg_length = 0
        running_reward = 0
        running_sr = 0
    episode_reward = 0

    if i_episode > args.num_episodes:
        break

# trainer.save_model(args.scenario, suffix="dads")
# print("Models saved to "+"models/{}/run{}/".format(args.scenario, args.run))
env.close()

now_time = datetime.datetime.now().strftime("%Y.%m.%d - %H:%M:%S")
exp_time = "{} -> {}".format(start_time, now_time)
print("This experiment time: {}".format(exp_time))

exp_utils.save_exp_config(args, exp_time, exp_name="dads-indie")


import json

import matplotlib.pyplot as plt
import networkx as nx
import pandas as pd
import scipy

def open_json(filename):
    with open(filename, "r", encoding="utf-8") as read_file:
        data = json.load(read_file)
    return data

class Person:
    def __init__(self, dictionary):
        self.id = dictionary['id']
        self.name = dictionary['name']
        self.shortname = dictionary['shortname']
        self.gender = dictionary['gender']
        if 'pids' in dictionary.keys():
            self.pids = dictionary['pids']
        else:
            self.pids = []
        if 'mid' in dictionary.keys():
            self.mid = dictionary['mid']
        else:
            self.mid = -1
        if 'fid' in dictionary.keys():
            self.fid = dictionary['fid']
        else:
            self.fid = -1
        self.cids = []
        if 'date' in dictionary.keys():
            self.date = dictionary['date']
        else:
            self.date = "-"
        if 'x' in dictionary.keys():
            self.x = dictionary['x']
        else:
            self.x = -1
        if 'y' in dictionary.keys():
            self.y = dictionary['y']
        else:
            self.y = -1
        self.rank = -1

    def __repr__(self):
        result = str(self.id) + ', ' + self.name + ', ' + self.gender
        result += ', ' + str(self.pids)
        result += ', ' + str(self.mid)
        result += ', ' + str(self.fid)
        result += ', ' + str(self.cids)
        result += ', ' + str(self.rank)
        return '{' + result + '}'

    def get_data(self):
        data = {'id': self.id,
                'pids': self.pids,
                'mid': self.mid,
                'fid': self.fid,
                'name': self.name,
                'gender': self.gender,
                'shortname': self.shortname,
                'date': self.date,
                'cids': self.cids,
                'rank': self.rank,
                'type': 'person',
                'x': self.x,
                'y': self.y}
        if data['mid'] == -1:
            data.pop('mid')
        if data['fid'] == -1:
            data.pop('fid')
        if data['x'] == -1:
            data.pop('x')
        if data['y'] == -1:
            data.pop('y')
        return data

class Family:
    def __init__(self, id, sid1, sid2, cids, x=-1, y=-1):
        self.id = id
        self.sid1 = sid1
        self.sid2 = sid2
        self.cids = cids
        self.x = x
        self.y = y

    def __repr__(self):
        result = str(self.id) + ', ' + str(self.sid1) + ', ' + str(self.sid2) + ', ' + str(self.cids)
        return '{' + result + '}'

    def __eq__(self, other):
        if ((self.sid1 == other.sid1 and self.sid2 == other.sid2) or (
                self.sid2 == other.sid1 and self.sid1 == other.sid2)) and self.cids == other.cids:
            return True
        else:
            return False

    def get_data(self):
        data = {'id': self.id,
                'type': 'family',
                'sid1': self.sid1,
                'sid2': self.sid2,
                'cids': self.cids,
                'x': self.x,
                'y': self.y}
        if data['x'] == -1:
            data.pop('x')
        if data['y'] == -1:
            data.pop('y')
        return data

class FamilyTree:
    def __init__(self, tree_list):
        self.tree = {}
        self.families = {}
        if type(tree_list) is list:
            for node in tree_list:
                if type(node) is dict:
                    if 'x' in tree_list[0].keys():
                        if node['type'] == 'person':
                            person = Person(node)
                            self.tree.update({person.id: person})
                        else:
                            family = Family(node['id'], node['sid1'], node['sid2'], node['cids'], node['x'], node['y'])
                            self.tree.update({family.id: family})
                            self.families.update({family.id: family})
                    else:
                        person = Person(node)
                        self.tree.update({person.id: person})
                else:
                    print("List must contains dict!")
        else:
            print("Tree must be list!")

        if 'x' not in tree_list[0].keys():
            self.__add_children()
            self.__add_rank()
            self.__add_families()

    def __add_children(self):
        result = self.tree
        for i in self.tree:
            if self.tree[i].mid != -1:
                result[self.tree[i].mid].cids.append(self.tree[i].id)
            if self.tree[i].fid != -1:
                result[self.tree[i].fid].cids.append(self.tree[i].id)
        self.tree = result

    def __add_rank(self):
        result = self.tree

        def add_rank_node(node, rank):
            if node.rank == -1:
                node.rank = str(rank)
                for n in node.pids:
                    add_rank_node(result[n], rank)
            else:
                rank += 1
                for n in node.cids:
                    add_rank_node(result[n], rank)

        add_rank_node(result[1], 0)
        self.tree = result

    def __add_families(self):
        i = 0
        for person in self.tree.values():
            for pid in person.pids:
                family = Family('f' + str(i), person.id, pid, list(set(person.cids) & set(self.tree[pid].cids)))
                if family not in self.families.values():
                    self.families.update({family.id: family})
                    i += 1

    def get_all_attr(self, attr_name):
        result = []
        for person in self.tree.values():
            result.append(getattr(person, attr_name))
        return result

    def get_dataframe(self):
        df_dict = {}
        for attr in ['id', 'name', 'gender', 'pids', 'mid', 'fid', 'cids', 'rank']:
            df_dict.update({attr: self.get_all_attr(attr)})
        df = pd.DataFrame.from_dict(df_dict)
        df.set_index('id', inplace=True)
        return df

    def get_persons_df(self):
        df_dict = {}
        for attr in ['id', 'name', 'gender', 'rank']:
            df_dict.update({attr: self.get_all_attr(attr)})
        df = pd.DataFrame.from_dict(df_dict)
        df.set_index('id', inplace=True)
        return df

    def get_links_df(self):
        df_dict = {}
        for attr in ['id', 'pids', 'cids']:
            df_dict.update({attr: self.get_all_attr(attr)})
        df = pd.DataFrame.from_dict(df_dict)
        df.set_index('id', inplace=True)
        return df

    def get_links_pair_df(self):
        df_dict = {'source_id': [], 'target_id': [], 'color': []}
        for person in self.tree.values():
            for pid in person.pids:
                if self.tree[pid].id not in df_dict['source_id'] and person.id not in df_dict['target_id']:
                    df_dict['source_id'].append(person.id)
                    df_dict['target_id'].append(self.tree[pid].id)
                    df_dict['color'].append('green')
            for cid in person.cids:
                df_dict['source_id'].append(person.id)
                df_dict['target_id'].append(self.tree[cid].id)
                df_dict['color'].append('blue')
        df = pd.DataFrame.from_dict(df_dict)
        return df

    def get_links_pair_families_df(self):
        df_dict = {'source_id': [], 'target_id': [], 'color': []}

        for family in self.families.values():
            df_dict['source_id'].append(family.sid1)
            df_dict['target_id'].append(family.id)
            df_dict['color'].append('green')
            df_dict['source_id'].append(family.sid2)
            df_dict['target_id'].append(family.id)
            df_dict['color'].append('green')
            for cid in family.cids:
                df_dict['source_id'].append(family.id)
                df_dict['target_id'].append(cid)
                df_dict['color'].append('blue')

        df = pd.DataFrame.from_dict(df_dict)
        return df

    def get_tree_family(self):
        tree_family = {}
        tree_family.update(self.tree)
        tree_family.update(self.families)
        return tree_family

    def get_networkx_graph(self):
        df = self.get_links_pair_families_df()
        G = nx.from_pandas_edgelist(df, 'source_id', 'target_id')
        for obj in self.get_tree_family().values():
            G.nodes[obj.get_data()['id']].update(obj.get_data())
        return G

# family_tree = FamilyTree(open_json('data/tree3_simple.json'))
# d = dict(family_tree.get_networkx_graph().nodes(data=True))
# print(d)
# print(list(d)[0])
# print(nx.adjacency_matrix(family_tree.get_networkx_graph()).toarray())