import cv2
import numpy as np
import matplotlib.pyplot as plt

img = cv2.imread('./test.jpg',cv2.WINDOW_NORMAL)
# cv2.namedWindow('transfrom')
rows,cols = img.shape

# 位移变换
M1 = np.float32([[1,0,50],[0,1,50]])
photo1 = cv2.warpAffine(img,M1,(rows,cols))
# cv2.imshow('transform',photo1)

# 旋转变换
M2 = cv2.getRotationMatrix2D((rows/2,cols/2),180,1)
photo2 = cv2.warpAffine(img,M2,(rows,cols))
# cv2.imshow('rotation',photo2)

# Affine 仿射变换(相对位置关系保留，输入三个点)
pts1 = np.float32([[50,50],[200,50],[50,200]])
pts2 = np.float32([[10,100],[100,50],[100,250]])
M3 = cv2.getAffineTransform(pts1,pts2)
photo3 = cv2.warpAffine(img,M3,(cols,rows))
#cv2.imshow('Affine',photo3)
plt.subplot(121),plt.imshow(img),plt.title('input')
plt.subplot(122),plt.imshow(photo3),plt.title('output')
plt.show()
cv2.waitKey(0)

# persipective 景深变换(输入四个点位置，任意三个非线性)
pts1 = np.float32([[56,65],[200,53],[28,200],[200,200]])
pts2 = np.float32([[0,0],[300,0],[0,300],[300,300]])
M4 = cv2.getPerspectiveTransform(pts1,pts2)
photo4 = cv2.warpPerspective(img,M4,(rows,cols))
plt.subplot(121),plt.imshow(img),plt.title('input')
plt.subplot(122),plt.imshow(photo4),plt.title('output')
plt.show()
cv2.destroyAllWindows()


import json
from calendar import monthrange
from dataclasses import dataclass, field
from datetime import date, datetime
from math import ceil
from typing import Any, NamedTuple
from zoneinfo import ZoneInfo

import typer
from rich import print
from rich.prompt import Prompt

from reportgen.utils.custom_types import CurrentAge, PatientData
from reportgen.utils.files import load_diagnostic, load_patient

def current_age(birthday: str) -> CurrentAge:
    start = datetime.strptime(birthday, "%Y-%m-%d %H:%M:%S").date()
    today = datetime.now(tz=ZoneInfo("America/Lima")).date()

    difference = today - start

    years = difference.days // 365
    months = (difference.days % 365) // 30
    days = (difference.days % 365) % 30

    return CurrentAge(
        years=abs(years),
        months=abs(months),
        days=abs(days),
        b=start,
        format=f"{start.day}     {start.month}     {start.year}",
        second_format=f"{start.day}/{start.month}/{start.year}",
    )

def input_patients(blocks: int) -> list[PatientData]:
    patients: list[PatientData] = []

    while True:
        print(f"\n[yellow]Queda(n) {blocks} espacio(s) en la hoja.[/yellow]")

        add_patient = typer.confirm("Agregar nuevo paciente?")
        if not add_patient:
            break

        dni = Prompt.ask("\n[blue]Número de DNI[/blue]")
        weight = Prompt.ask("[blue]Peso en kg[/blue]")
        size = Prompt.ask("[blue]Talla en cm[/blue]")
        hb = Prompt.ask("[blue]Valor de Hb[/blue]")

        patient = load_patient(dni=dni)
        age = current_age(birthday=patient.birthday)
        his = load_diagnostic(age=age)

        patients.append(
            PatientData(
                personal={
                    "dni": dni,
                    "weight": weight,
                    "size": size,
                    "hb": hb,
                },
                identification=patient,
                his=his,
                age=age,
            )
        )

        blocks -= ceil(len(his) / 3)

    return patients

GENDER = {0: "NO ESPECIFICADO", 1: "MASCULINO", 2: "FEMENINO"}

INSURANCE = {
    0: "NO ESPECIFICADO",
    1: "SIS",
    2: "ESSALUD",
    3: "OTRO",
}

class Age(NamedTuple):
    years: int
    months: int
    days: int

    def __str__(self) -> str:
        return (
            f"{self.years} año(s), {self.months} mes(es) y {self.days} día(s)"
        )

class Control(NamedTuple):
    cie: str
    description: str
    dx: str
    lab: tuple[str, str, str]

@dataclass
class TempPatient:
    dni: str
    hlc: str = field(init=False)
    district: str = field(init=False)
    establishment: str = field(init=False)
    sector: str = field(init=False)
    insurance: str = field(init=False)
    full_name: str = field(init=False)
    gender: str = field(init=False)
    birthday: date = field(init=False)
    age: Age = field(init=False)
    appointment: str = field(init=False)
    type_of_birth: int = field(init=False)
    diagnostic: list[Control] = field(init=False)
    age_diagnostic: Age = field(init=False)

    def __post_init__(self) -> None:
        patient = self.load_from_json()

        self.hlc = patient["hlc"]
        self.district = patient["district"]
        self.establishment = patient["establishment"]
        self.sector = patient["sector"]
        self.insurance = INSURANCE[int(patient["insurance"])]
        self.full_name = f"{patient['father_last_name']} {patient['mother_last_name']}, {patient['names']}"
        self.gender = GENDER[int(patient["gender"])]
        self.birthday = datetime.strptime(
            patient["birthday"], "%Y-%m-%d %H:%M:%S"
        ).date()
        self.age = self.current_age()
        self.appointment = patient["appointment"]
        self.type_of_birth = patient["type_of_birth"]
        self.diagnostic = self.load_diagnostic()

    def load_from_json(self) -> Any:
        with open("database/people.json", encoding="utf-8") as f:
            patients = json.load(f)

        return patients[f"{self.dni}"]

    def current_age(self) -> Age:
        year, month, day, *_ = datetime.now(
            tz=ZoneInfo("America/Lima")
        ).timetuple()

        if day < self.birthday.day:
            month -= 1
            day += monthrange(self.birthday.year, self.birthday.month)[1]

        if month < self.birthday.month:
            year -= 1
            month += 12

        return Age(
            year - self.birthday.year,
            month - self.birthday.month,
            day - self.birthday.day,
        )

    def load_diagnostic(self) -> list[Control]:
        with open("database/codes.json", encoding="utf-8") as f:
            diagnostics = json.load(f)

        # h = Age(0, 0, 0)

        try:
            if self.age.years == 0 and self.age.months == 0:
                raw_diagnostic = diagnostics["RN"][f"{self.age.days}_days"]
            elif self.age.years > 5:
                raw_diagnostic = diagnostics["5_years"][
                    f"{self.age.months}_months"
                ]
            else:
                raw_diagnostic = diagnostics[f"{self.age.years}_years"][
                    f"{self.age.months}_months"
                ]
        except KeyError:
            # for i in range(1, 8):
            # print(i)
            raw_diagnostic = []
            # print(key)

        return [
            Control(
                cie=d[0],
                description=d[1],
                dx=d[2],
                lab=tuple(d[3]),  # type: ignore
            )
            for d in raw_diagnostic
        ]


from os.path import join

import joblib
import pandas as pd
from sklearn.tree import export_graphviz
from sklearn.utils import resample

from config import DIR_PATH, PATH_MATRICES
from disambiguator import Disambiguator, TARGET, CATEGORICAL_FEATURES
from utils.helpers import extract_file_name_from_path, convert_str_to_list, load_pickle_file

def convert_binary_to_triple_schema(assignment):
    if assignment is True:
        return 1
    elif assignment is False:
        return -1
    else:
        return 0

def downsample(x, y, max_skewness=1, max_size=None, return_index=False):
    """Adopted from https://elitedatascience.com/imbalanced-classes"""
    df = x.copy()
    df.loc[:, 'TARGET'] = y
    df_class_1 = df[(df['TARGET'] == "true") | (df['TARGET'] == 1)]
    df_class_0 = df[(df['TARGET'] == "false") | (df['TARGET'] == 0)]

    # Downsample class 1
    if max_skewness == 0:  # change nothing
        size = len(df_class_1)
    elif max_skewness >= 1:
        size = int(min(max_skewness * len(df_class_0), len(df_class_1)))
    else:  # max_skewness < 1 but > 0, then increase reduce size of class 1 even more so that class 0 is overrepresented
        size = int(max_skewness * len(df_class_0))
    try:
        df_class_1_downsampled = resample(df_class_1, replace=False, n_samples=size, random_state=1)
    except ValueError:
        print("Class 0 still too large, smapling class 1 with replacement")
        df_class_1_downsampled = resample(df_class_1, replace=True, n_samples=size, random_state=1)

    # Combine minority class with downsampled majority class
    df_downsampled = pd.concat([df_class_1_downsampled, df_class_0])
    if max_size and max_size < len(df_downsampled):
        df_downsampled = df_downsampled.sample(max_size, random_state=1)

    if return_index is True:
        return df_downsampled.index
    else:
        return df_downsampled.drop('TARGET', axis=1), df_downsampled['TARGET']

def export_decisiontree(tree_model, file_path, feature_names):
    export_graphviz(tree_model, out_file=file_path, precision=2,
                    filled=True, feature_names=feature_names)

def load_ai_blocks(files, max_size, min_size):
    blocks = []
    print(len(files))
    for f in files:
        df = pd.read_excel(f)
        if max_size >= len(df) > min_size:
            blocks.append(extract_file_name_from_path(f)[0])
    return blocks

def build_similarity_matrices(blocks, features, verbose=False, squared=False):
    """for each ai in blocks, build a DataFrame containing the attributes and a DataFrame containing all
     pairwise combinations (=features) of all authorships in the block.
     Return dictionaries {ai: attributes_df} and {ai: features_df}"""

    def build_similarity_matrix(block, features, verbose=False, squared=False):
        d = Disambiguator(block)
        """get attrs from features"""
        feature_to_attr = d.feature_to_pw_comp_method
        create_attrs = [v[0] for f, v in feature_to_attr.items() if f in features and f != "aid_equal"]
        create_attrs = list(set(create_attrs))
        print(f"create attrs: {create_attrs}")
        d.populate_block(verbose=verbose)
        for col in ['country', 'city']:  # entries are list objects but stored in file as strings
            if col in d.block.columns:
                d.block.loc[:, col] = d.block.loc[:, col].map(convert_str_to_list)
        d.ignore_ambiguous_authorships()
        d.prepare_attributes_for_modelling(create_attrs=create_attrs)
        if verbose:
            print("Attributes: ", d.attributes.columns)

        d.build_all_pw_comparisons(store_pw_combs=True, squared=squared, features=features)
        feats = d.feature_matrix.merge(d.pw_combs[['document_id_x', 'document_id_y']], how='outer',
                                       left_index=True,
                                       right_index=True)  # TODO: this simply adds document_ids and should be replaced by real IDs

        return feats, d.attributes

    print(f"features: {features}")
    ai_block_to_similarity_matrix = {}
    ai_block_to_all_attrs = {}
    print('...')
    print(len(blocks))

    for ai_block in blocks:
        if verbose:
            print(ai_block)
        feats, attrs = build_similarity_matrix(ai_block, features, verbose, squared)

        ai_block_to_all_attrs[ai_block] = attrs
        ai_block_to_similarity_matrix[ai_block] = feats

    return ai_block_to_similarity_matrix, ai_block_to_all_attrs

def remove_document_id_col(df):
    cols = [c for c in df.columns if not c.startswith('document_id')]
    return df[cols].copy()

def remove_duplicate_rows(df):
    return df.drop_duplicates().reset_index(drop=True)

def split_into_features_and_target(df, target_col):
    d = df.copy()
    d.loc[:, target_col] = d[target_col].map({'true': 1, 'false': 0})
    features_df, target = d.drop(target_col, axis=1), d[target_col]
    return features_df, target

def vectorize_categorical_features(df, target_col, cat_attrs):
    cat_cols = [c for c in df.columns if c in cat_attrs and c != target_col]
    return pd.get_dummies(df, columns=cat_cols)

def encode_path_params_for_model_dump(model_type, scoring_func, max_block_size, min_block_size, downsample_ratio):
    file_name_part = '_'.join([model_type, scoring_func, 'blocksize_range', str(min_block_size), str(max_block_size),
                               'max_skewness', str(downsample_ratio)])
    file_name_clf = file_name_part + '.pkl'
    file_name_validation_blocks = 'ai_blocks_validate_' + file_name_part + '.pkl'
    file_name_feature_info = 'feature_info_' + file_name_part + '.pkl'

    return file_name_clf, file_name_validation_blocks, file_name_feature_info

def load_model(model_type, scoring_func, max_block_size, min_block_size, downsample_ratio):
    file_name_clf, file_name_validation_blocks, \
    file_name_feature_info = encode_path_params_for_model_dump(model_type, scoring_func, max_block_size, min_block_size,
                                                               downsample_ratio)
    clf = joblib.load(join(DIR_PATH, 'models', file_name_clf))
    feature_info = joblib.load(join(DIR_PATH, 'models', file_name_feature_info))

    return clf, feature_info  # , ai_blocks_validate

def remove_obvious_nonmatches(df):
    """Remove rows from feature matrix (each row represents comparison of two authorships) where
    first_name or middle_name or middle_name_init are incompatible"""
    cond_first_name = (df['first_name_equal'] != 'false')
    cond_middle_name = (df['middle_name_equal'] != 'false')
    cond_middle_name_init = (df['middle_name_init_equal'] != 'false')

    return df[cond_first_name & cond_middle_name & cond_middle_name_init]

def filter_by_ai_blocks(block_to_f_mat, block_to_a_mat, blocks):
    block_to_f_mat_reduced = {k: v for k, v in block_to_f_mat.items() if k in blocks}
    block_to_a_mat_reduced = {k: v for k, v in block_to_a_mat.items() if k in blocks}

    return blocks, block_to_f_mat_reduced, block_to_a_mat_reduced

def fill_missing_feature_cols(df, feature_cols):
    for f in set(feature_cols).difference(set(df.columns)):
        df[f] = 0
    df = df.reindex(columns=feature_cols)
    return df

def filter_ai_blocks_by_size(block_to_matrix, block_to_attr, max_size, min_size):
    blocks = [b for b, m in block_to_attr.items() if max_size >= len(m) > min_size]
    bm = {b: m for b, m in block_to_matrix.items() if b in blocks}
    ba = {b: a for b, a in block_to_attr.items() if b in blocks}

    return blocks, bm, ba

def remove_au_pairs_with_same_doc_id(feat_matrix):
    doc_id_cols = [c for c in feat_matrix.columns if c.startswith('document_id')]
    assert len(doc_id_cols) == 2
    col_1, col_2 = doc_id_cols[0], doc_id_cols[1]
    """
    document_id cols contain author position after + sign.
    Restrict to only those rows (authorship pairs) where doc id is different.
    """
    return feat_matrix[
        feat_matrix[col_1].map(lambda x: x.split('+')[0]) != feat_matrix[col_2].map(
            lambda x: x.split('+')[0])].reset_index(drop=True)

def prepare_sim_matrix_for_clustering(feat_matrix, features):
    """Does not reduce the x-dimension of matrix, only manipulation of features (columns)"""
    data_cluster = remove_document_id_col(feat_matrix)
    x = split_into_features_and_target(data_cluster, TARGET)[0]
    x = vectorize_categorical_features(x, TARGET, CATEGORICAL_FEATURES)
    x = fill_missing_feature_cols(x, features)
    return x

def load_matrices(blocks):
    ai_block_to_matrix = load_pickle_file(join(PATH_MATRICES, 'ai_block_to_matrix.pickle'))
    ai_block_to_attrs = load_pickle_file(join(PATH_MATRICES, 'ai_block_to_attrs.pickle'))
    ai_blocks, ai_block_to_matrix, ai_block_to_attrs = filter_by_ai_blocks(ai_block_to_matrix, ai_block_to_attrs,
                                                                           blocks)

    return ai_block_to_attrs, ai_block_to_matrix


from detection_msgs.msg import BoundingBoxes
import numpy as np
import rospy as ros
from sensor_msgs import point_cloud2
from sensor_msgs.msg import PointCloud2
from robotica.msg import Coord
import open3d as o3d
import time

class vision():

    def __init__(self):
        ## Matrice di rotazione da camera frame a world frame
        self.w_R_c = np.array([[0.0, -0.49948, 0.86632],
                               [-1.0, 0.0, 0.0],
                              [-0.0, -0.86632, -0.49948]])
        ## Vettore di traslazione da robot frame a camera frame
        self.x_c = np.array([-0.4, 0.59, 1.4])

        self.point_cloud = PointCloud2()
        self.voxel_size = 0.003
        self.models="/home/lucaboschiero/ros_ws/src/locosim/ros_impedance_controller/worlds/models"

    def pointcloud_callback(self, msg):
        # Salva la pointcloud della camera
        self.point_cloud = msg

    def pointCloud_fpfh(self, pointCloud):

        #La pointcloud viene sottocampionata
        pointCloud_down = o3d.geometry.PointCloud.voxel_down_sample(pointCloud, self.voxel_size)

        # Calcolare le normali per ogni punto nella point cloud ridotta
        radius = self.voxel_size * 2
        o3d.geometry.PointCloud.estimate_normals(
            pointCloud_down,
            o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=30))

        radius = self.voxel_size * 5
        # Calcola le feature FPFH sulla point cloud ridotta
        pointCloud_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
            pointCloud_down,
            o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=100))

        return pointCloud_down, pointCloud_fpfh

    def registration(self, inputPointCloud_down, stlPointCloud_down, inputPointCloud_fpfh, stlPointCloud_fpfh):
        # Esegue la registrazione tra due point cloud utilizzando Open3D

        # Esegue la registrazione RANSAC basata sulla corrispondenza delle feature
        distance_threshold = self.voxel_size * 1.5

        result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
            inputPointCloud_down, stlPointCloud_down, inputPointCloud_fpfh, stlPointCloud_fpfh, True, distance_threshold,
            o3d.pipelines.registration.TransformationEstimationPointToPoint(False), 4, [
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(
                    distance_threshold)
            ], o3d.pipelines.registration.RANSACConvergenceCriteria(1000000, 500))

        distance_threshold = self.voxel_size * 0.4
        result1 = o3d.pipelines.registration.registration_icp(
            inputPointCloud_down, stlPointCloud_down, distance_threshold, result.transformation,
            o3d.pipelines.registration.TransformationEstimationPointToPlane())

        return result1

    def getModelPath(self, cl):

        # In base alla classe del blocco ritorna il path del modello stl

        if cl == "1" :
            return "/model0/meshes/X1-Y1-Z2.stl"

        elif cl == "2" :
            return "/model1/meshes/X1-Y2-Z1.stl"

        elif cl == "3" :
            return "/model2/meshes/X1-Y2-Z2.stl"

        elif cl == "4" :
            return "/model3/meshes/X1-Y2-Z2-CHAMFER.stl"

        elif cl == "5" :
            return "/model4/meshes/X1-Y2-Z2-TWINFILLET.stl"

        elif cl == "6" :
            return "/model5/meshes/X1-Y3-Z2.stl"

        elif cl == "6-FILLER" :
            return "/model6/meshes/X1-Y3-Z2-FILLET.stl"

        elif cl == "7" :
            return "/model7/meshes/X1-Y4-Z1.stl"

        elif cl == "8" :
            return "/model8/meshes/X1-Y4-Z2.stl"

        elif cl == "9" :
            return "/model9/meshes/X2-Y2-Z2.stl"

        elif cl == "9-FILLER" :
            return "/model10/meshes/X2-Y2-Z2-FILLET.stl"

    def pose_estimation(self, points, object_class):
        #Trasforma i punti della pointcloud di input in un vettore open3d
        points = np.asarray(points)
        inputPointCloud = o3d.geometry.PointCloud()
        inputPointCloud.points = o3d.utility.Vector3dVector(points)

        # Ottiene il path in base al tipo di blocco
        model_path = self.getModelPath(object_class)
        # Legge il file .stl
        stl_file_path = self.models + model_path
        try:
            print(stl_file_path)
            mesh = o3d.io.read_triangle_mesh(stl_file_path)
        except Exception as e:
            print(f"Error reading STL file: {str(e)}")
            return None

        # Crea una pointcloud a partire da una mesh tridimensionale
        nPoints = int (np.size(points))
        mesh_points = mesh.sample_points_uniformly(number_of_points=nPoints)
        stlPointCloud = o3d.geometry.PointCloud()
        stlPointCloud.points = mesh_points.points

        # Sottodimensiona le due pointcloud e estrae le fpfh features
        inputPointCloud_down, inputPointCloud_fpfh = self.pointCloud_fpfh(inputPointCloud)
        stlPointCloud_down, stlPointCloud_fpfh = self.pointCloud_fpfh(stlPointCloud)

        # Esegue la registrazione delle pointcloud
        result = self.registration(inputPointCloud_down, stlPointCloud_down, inputPointCloud_fpfh, stlPointCloud_fpfh)

        # Ottiene la matrice di rotazione
        R = result.transformation[:3, :3]

        return R

    def detection_callback(self, msg):

        # Topic dove pubblica le coordinate dei blocchi e la classe
        pubco = ros.Publisher("/coordinates", Coord, queue_size=100)

        # Crea un messaggio dove salvare i dati da pubblicare
        coordinate = Coord()

        print("Detection time: ",msg.time)

        #salva il tempo di inizio
        start_time = time.time()

        # Scorre tra gli elementi trovati da Yolo
        for i in range(msg.count):
            points_list = []

            x_min = msg.bounding_boxes[i].xmin
            x_max = msg.bounding_boxes[i].xmax
            y_min = msg.bounding_boxes[i].ymin
            y_max = msg.bounding_boxes[i].ymax

            # Riceve i punti della bounding box del blocco e calcola u e v
            u = int((msg.bounding_boxes[i].xmin + msg.bounding_boxes[i].xmax) / 2)
            v = int((msg.bounding_boxes[i].ymin + msg.bounding_boxes[i].ymax) / 2)
            cl = (msg.bounding_boxes[i].Class)

            # Trasnforma u e v in coordinate del camera frame
            for data in point_cloud2.read_points(self.point_cloud, field_names=['x', 'y', 'z'], skip_nans=False, uvs=[(u, v)]):
                points_list.append([data[0], data[1], data[2]])
                print("Data Optical frame: ", points_list)

            # Trasforma le coordinate nel world frame
            pointW = self.w_R_c.dot(points_list[0]) + self.x_c

            # Salva x, y, z e la classe del blocco per poi pubblicarle
            coordinate.x = pointW[0]
            coordinate.y = pointW[1]
            coordinate.z = pointW[2]
            coordinate.cl = cl

            # Genera un vettore di punti
            uv = [(x, y) for y in range(y_min, y_max + 1) for x in range(x_min, x_max + 1)]

            # Trasforma i punti in coordinate x, y e z rispetto al camera frame
            point_region = []
            for data in point_cloud2.read_points(self.point_cloud, field_names=['x', 'y', 'z'], skip_nans=False, uvs=uv):
                point_region.append([data[0], data[1], data[2]])

            # Trasforma i punti dal frame della camera al world frame
            points = []
            for point in point_region:
                points.append(self.w_R_c.dot(point) + self.x_c)

            # Ottiene il centro dei punti
            pointMiddle = np.mean(points, axis=0)

            # Applica l'offset di tutti i punti dal centro
            pointsOffset = []
            for point in points:
                pointsOffset.append((point - pointMiddle))

            # Ottiene la matrice di rotazione
            rotation_matrix = self.pose_estimation(pointsOffset, cl)
            print("Rotation matrix: ",rotation_matrix)

            # Inserisce nel messaggio la matrice di rotazione
            coordinate.R = rotation_matrix.flatten().tolist()

            end_time = time.time()
            # Calcola il tempo trascorso in secondi
            total_time = end_time - start_time

            print("Orientation detection time:", total_time, "secondi")

            # Pubblica il messaggio
            pubco.publish(coordinate)

def main(p):

    # Vengono definiti i topic a cui è sottoscritto il codice
    ros.init_node("listener")

    ros.Subscriber("/yolov5/detections", BoundingBoxes, p.detection_callback)
    ros.Subscriber("/ur5/zed_node/point_cloud/cloud_registered", PointCloud2, p.pointcloud_callback)

    ros.spin()

if __name__ == '__main__':

    p = vision()
    try:
        main(p)
    except ros.ROSInterruptException:
        pass


import torch
import torch.nn as nn

def double_conv(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels, out_channels, 3, padding=1),
        nn.ReLU(inplace=True))

class UNet(nn.Module):

    def __init__(self, n_classes):
        super().__init__()

        self.conv_down1 = double_conv(3, 64)
        self.conv_down2 = double_conv(64, 128)
        self.conv_down3 = double_conv(128, 256)
        self.conv_down4 = double_conv(256, 512)

        self.maxpool = nn.MaxPool2d(2)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv_up3 = double_conv(256 + 512, 256)
        self.conv_up2 = double_conv(128 + 256, 128)
        self.conv_up1 = double_conv(128 + 64, 64)

        self.last_conv = nn.Conv2d(64, n_classes, kernel_size=1)

    def forward(self, x):
        # Batch - 1d tensor.  N_channels - 1d tensor, IMG_SIZE - 2d tensor.
        # Example: x.shape >>> (10, 3, 256, 256).

        conv1 = self.conv_down1(x)
        x = self.maxpool(conv1)
        conv2 = self.conv_down2(x)
        x = self.maxpool(conv2)
        conv3 = self.conv_down3(x)
        x = self.maxpool(conv3)
        x = self.conv_down4(x)
        x = self.upsample(x)

        x = torch.cat([x, conv3], dim=1)

        x = self.conv_up3(x)
        x = self.upsample(x)
        x = torch.cat([x, conv2], dim=1)

        x = self.conv_up2(x)
        x = self.upsample(x)
        x = torch.cat([x, conv1], dim=1)

        x = self.conv_up1(x)

        out = self.last_conv(x)
        out = torch.sigmoid(out)

        return out


from selenium import webdriver
from selenium.webdriver.common.by import By

driver = webdriver.Chrome()
driver.get('https://m.place.naver.com/restaurant/1239997161/home')

images_container = driver.find_element(By.CLASS_NAME, "CB8aP")
images = images_container.find_elements(By.CLASS_NAME, "K0PDV")

for image in images:
    print(image.value_of_css_property("background-image").split("\"")[1])
driver.quit()


import sys
import os
import json
import requests
from flask import Flask, request, abort
import tokens
GITHUB_SECRET = tokens.github
bearer = tokens.webex
roomId = tokens.roomId
headers = {
    "Accept": "application/json",
    "Content-Type": "application/json; charset=utf-8",
    "Authorization": "Bearer " + bearer
}

#utils library

def send_get(url, payload=None,js=True):
    if payload == None:
        request = requests.get(url, headers=headers)
    else:
        request = requests.get(url, headers=headers, params=payload)
    if js == True:
        request= request.json()
    return request

def send_post(url, data):
    request = requests.post(url, json.dumps(data), headers=headers).json()
    return request

def help_me():
    return "Hello! Sure! I can help. Below are the commands that I understand: `Help me` - I will display what I can do. My job is to send you and the rest of your team notifications when something happens inside your repository on pull or push events"

def greetings():
    return "Hi, type `Help me` to see what I can do.<br/>"

def send_msg(data):
    if str(data['action']) == 'opened':
        msg = 'the user ' +str(data['pull_request']['user']['login'])+' opened a pull request at url: '+str(data['pull_request']['url'])
        send_post("https://api.ciscospark.com/v1/messages", {"roomId": roomId,"markdown": msg})
    elif str(data['action']) == 'unlabeled':
        msg = 'the user ' +str(data['pull_request']['user']['login'])+' made a request to unlabel at url: '+str(data['pull_request']['url'])
        send_post("https://api.ciscospark.com/v1/messages", {"roomId": roomId,"markdown": msg})
    elif str(data['action']) == 'labeled':
        msg = 'the user ' +str(data['pull_request']['user']['login'])+' made a request to label at url: '+str(data['pull_request']['url'])
        send_post("https://api.ciscospark.com/v1/messages", {"roomId": roomId,"markdown": msg})
    elif str(data['action']) == 'closed':
        msg = 'the user ' +str(data['pull_request']['user']['login'])+' closed a pull request at url: '+str(data['pull_request']['url'])
        send_post("https://api.ciscospark.com/v1/messages", {"roomId": roomId,"markdown": msg})
    else:
        msg = 'the user ' +str(data['pull_request']['user']['login'])+' made a request at url: '+str(data['pull_request']['url'])
        send_post("https://api.ciscospark.com/v1/messages", {"roomId": roomId,"markdown": msg})


import cv2 as cv
from collections import defaultdict
cvNet = cv.dnn.readNetFromTensorflow('/Users/Tinku/Desktop/RoboticsShit/Skystone-Vision/Final/image_tensor/PB_file/model.pb', '/Users/Tinku/Desktop/RoboticsShit/Skystone-Vision/Final/image_tensor/Pbtxt/model.pbtxt')

img = cv.imread('/Users/Tinku/Desktop/TestStuff/work.jpeg')
rows = img.shape[0]
cols = img.shape[1]
cvNet.setInput(cv.dnn.blobFromImage(img, size=(300, 300), swapRB=True, crop=False))
cvOut = cvNet.forward()
numOfRect = 0
objects_dict = {}
for detection in cvOut[0,0,:,:]:
    score = float(detection[2])
    objects = detection
    if score > 0.75:
        left = detection[3] * cols
        top = detection[4] * rows
        right = detection[5] * cols
        bottom = detection[6] * rows
        objects_dict[score]=[left,top,right,bottom]
        numOfRect += 1

        print(score)
#print(objects_dict)
l=list(objects_dict.keys())
l.sort(reverse=True)
print(l)
for i in l[:6] :
    print(i,objects_dict[i])
    cv.rectangle(img, (int(objects_dict[i][0]), int(objects_dict[i][1])), (int(objects_dict[i][2]), int(objects_dict[i][3])), (23, 230, 210), thickness=5)

cv.imshow('img', img)
print(numOfRect)
print(objects)
cv.waitKey()


"""
State Space Representation

Author: Chad Fulton
License: Simplified-BSD
"""

import warnings
import numpy as np
from .tools import (
    find_best_blas_type, validate_matrix_shape, validate_vector_shape
)
from .initialization import Initialization
from . import tools

class OptionWrapper:
    def __init__(self, mask_attribute, mask_value):
        # Name of the class-level bitmask attribute
        self.mask_attribute = mask_attribute
        # Value of this option
        self.mask_value = mask_value

    def __get__(self, obj, objtype):
        # Return True / False based on whether the bit is set in the bitmask
        return bool(getattr(obj, self.mask_attribute, 0) & self.mask_value)

    def __set__(self, obj, value):
        mask_attribute_value = getattr(obj, self.mask_attribute, 0)
        if bool(value):
            value = mask_attribute_value | self.mask_value
        else:
            value = mask_attribute_value & ~self.mask_value
        setattr(obj, self.mask_attribute, value)

class MatrixWrapper:
    def __init__(self, name, attribute):
        self.name = name
        self.attribute = attribute
        self._attribute = '_' + attribute

    def __get__(self, obj, objtype):
        matrix = getattr(obj, self._attribute, None)
        # # Remove last dimension if the array is not actually time-varying
        # if matrix is not None and matrix.shape[-1] == 1:
        #     return np.squeeze(matrix, -1)
        return matrix

    def __set__(self, obj, value):
        value = np.asarray(value, order="F")
        shape = obj.shapes[self.attribute]

        if len(shape) == 3:
            value = self._set_matrix(obj, value, shape)
        else:
            value = self._set_vector(obj, value, shape)

        setattr(obj, self._attribute, value)
        obj.shapes[self.attribute] = value.shape

    def _set_matrix(self, obj, value, shape):
        # Expand 1-dimensional array if possible
        if (value.ndim == 1 and shape[0] == 1 and
                value.shape[0] == shape[1]):
            value = value[None, :]

        # Enforce that the matrix is appropriate size
        validate_matrix_shape(
            self.name, value.shape, shape[0], shape[1], obj.nobs
        )

        # Expand time-invariant matrix
        if value.ndim == 2:
            value = np.array(value[:, :, None], order="F")

        return value

    def _set_vector(self, obj, value, shape):
        # Enforce that the vector has appropriate length
        validate_vector_shape(
            self.name, value.shape, shape[0], obj.nobs
        )

        # Expand the time-invariant vector
        if value.ndim == 1:
            value = np.array(value[:, None], order="F")

        return value

class Representation:
    r"""
    State space representation of a time series process

    Parameters
    ----------
    k_endog : {array_like, int}
        The observed time-series process :math:`y` if array like or the
        number of variables in the process if an integer.
    k_states : int
        The dimension of the unobserved state process.
    k_posdef : int, optional
        The dimension of a guaranteed positive definite covariance matrix
        describing the shocks in the measurement equation. Must be less than
        or equal to `k_states`. Default is `k_states`.
    initial_variance : float, optional
        Initial variance used when approximate diffuse initialization is
        specified. Default is 1e6.
    initialization : Initialization object or str, optional
        Initialization method for the initial state. If a string, must be one
        of {'diffuse', 'approximate_diffuse', 'stationary', 'known'}.
    initial_state : array_like, optional
        If `initialization='known'` is used, the mean of the initial state's
        distribution.
    initial_state_cov : array_like, optional
        If `initialization='known'` is used, the covariance matrix of the
        initial state's distribution.
    nobs : int, optional
        If an endogenous vector is not given (i.e. `k_endog` is an integer),
        the number of observations can optionally be specified. If not
        specified, they will be set to zero until data is bound to the model.
    dtype : np.dtype, optional
        If an endogenous vector is not given (i.e. `k_endog` is an integer),
        the default datatype of the state space matrices can optionally be
        specified. Default is `np.float64`.
    design : array_like, optional
        The design matrix, :math:`Z`. Default is set to zeros.
    obs_intercept : array_like, optional
        The intercept for the observation equation, :math:`d`. Default is set
        to zeros.
    obs_cov : array_like, optional
        The covariance matrix for the observation equation :math:`H`. Default
        is set to zeros.
    transition : array_like, optional
        The transition matrix, :math:`T`. Default is set to zeros.
    state_intercept : array_like, optional
        The intercept for the transition equation, :math:`c`. Default is set to
        zeros.
    selection : array_like, optional
        The selection matrix, :math:`R`. Default is set to zeros.
    state_cov : array_like, optional
        The covariance matrix for the state equation :math:`Q`. Default is set
        to zeros.
    **kwargs
        Additional keyword arguments. Not used directly. It is present to
        improve compatibility with subclasses, so that they can use `**kwargs`
        to specify any default state space matrices (e.g. `design`) without
        having to clean out any other keyword arguments they might have been
        passed.

    Attributes
    ----------
    nobs : int
        The number of observations.
    k_endog : int
        The dimension of the observation series.
    k_states : int
        The dimension of the unobserved state process.
    k_posdef : int
        The dimension of a guaranteed positive
        definite covariance matrix describing
        the shocks in the measurement equation.
    shapes : dictionary of name:tuple
        A dictionary recording the initial shapes
        of each of the representation matrices as
        tuples.
    initialization : str
        Kalman filter initialization method. Default is unset.
    initial_variance : float
        Initial variance for approximate diffuse
        initialization. Default is 1e6.

    Notes
    -----
    A general state space model is of the form

    .. math::

        y_t & = Z_t \alpha_t + d_t + \varepsilon_t \\
        \alpha_t & = T_t \alpha_{t-1} + c_t + R_t \eta_t \\

    where :math:`y_t` refers to the observation vector at time :math:`t`,
    :math:`\alpha_t` refers to the (unobserved) state vector at time
    :math:`t`, and where the irregular components are defined as

    .. math::

        \varepsilon_t \sim N(0, H_t) \\
        \eta_t \sim N(0, Q_t) \\

    The remaining variables (:math:`Z_t, d_t, H_t, T_t, c_t, R_t, Q_t`) in the
    equations are matrices describing the process. Their variable names and
    dimensions are as follows

    Z : `design`          :math:`(k\_endog \times k\_states \times nobs)`

    d : `obs_intercept`   :math:`(k\_endog \times nobs)`

    H : `obs_cov`         :math:`(k\_endog \times k\_endog \times nobs)`

    T : `transition`      :math:`(k\_states \times k\_states \times nobs)`

    c : `state_intercept` :math:`(k\_states \times nobs)`

    R : `selection`       :math:`(k\_states \times k\_posdef \times nobs)`

    Q : `state_cov`       :math:`(k\_posdef \times k\_posdef \times nobs)`

    In the case that one of the matrices is time-invariant (so that, for
    example, :math:`Z_t = Z_{t+1} ~ \forall ~ t`), its last dimension may
    be of size :math:`1` rather than size `nobs`.

    References
    ----------
    .. [*] Durbin, James, and Siem Jan Koopman. 2012.
       Time Series Analysis by State Space Methods: Second Edition.
       Oxford University Press.
    """

    endog = None
    r"""
    (array) The observation vector, alias for `obs`.
    """
    design = MatrixWrapper('design', 'design')
    r"""
    (array) Design matrix: :math:`Z~(k\_endog \times k\_states \times nobs)`
    """
    obs_intercept = MatrixWrapper('observation intercept', 'obs_intercept')
    r"""
    (array) Observation intercept: :math:`d~(k\_endog \times nobs)`
    """
    obs_cov = MatrixWrapper('observation covariance matrix', 'obs_cov')
    r"""
    (array) Observation covariance matrix:
    :math:`H~(k\_endog \times k\_endog \times nobs)`
    """
    transition = MatrixWrapper('transition', 'transition')
    r"""
    (array) Transition matrix:
    :math:`T~(k\_states \times k\_states \times nobs)`
    """
    state_intercept = MatrixWrapper('state intercept', 'state_intercept')
    r"""
    (array) State intercept: :math:`c~(k\_states \times nobs)`
    """
    selection = MatrixWrapper('selection', 'selection')
    r"""
    (array) Selection matrix:
    :math:`R~(k\_states \times k\_posdef \times nobs)`
    """
    state_cov = MatrixWrapper('state covariance matrix', 'state_cov')
    r"""
    (array) State covariance matrix:
    :math:`Q~(k\_posdef \times k\_posdef \times nobs)`
    """

    def __init__(self, k_endog, k_states, k_posdef=None,
                 initial_variance=1e6, nobs=0, dtype=np.float64,
                 design=None, obs_intercept=None, obs_cov=None,
                 transition=None, state_intercept=None, selection=None,
                 state_cov=None, statespace_classes=None, **kwargs):
        self.shapes = {}

        # Check if k_endog is actually the endog array
        endog = None
        if isinstance(k_endog, np.ndarray):
            endog = k_endog
            # If so, assume that it is either column-ordered and in wide format
            # or row-ordered and in long format
            if (endog.flags['C_CONTIGUOUS'] and
                    (endog.shape[0] > 1 or nobs == 1)):
                endog = endog.T
            k_endog = endog.shape[0]

        # Endogenous array, dimensions, dtype
        self.k_endog = k_endog
        if k_endog < 1:
            raise ValueError('Number of endogenous variables in statespace'
                             ' model must be a positive number.')
        self.nobs = nobs

        # Get dimensions from transition equation
        if k_states < 1:
            raise ValueError('Number of states in statespace model must be a'
                             ' positive number.')
        self.k_states = k_states
        self.k_posdef = k_posdef if k_posdef is not None else k_states

        # Make sure k_posdef <= k_states
        # TODO: we could technically allow k_posdef > k_states, but the Cython
        # code needs to be more thoroughly checked to avoid seg faults.
        if self.k_posdef > self.k_states:
            raise ValueError('Dimension of state innovation `k_posdef` cannot'
                             ' be larger than the dimension of the state.')

        # Bind endog, if it was given
        if endog is not None:
            self.bind(endog)

        # Record the shapes of all of our matrices
        # Note: these are time-invariant shapes; in practice the last dimension
        # may also be `self.nobs` for any or all of these.
        self.shapes = {
            'obs': (self.k_endog, self.nobs),
            'design': (self.k_endog, self.k_states, 1),
            'obs_intercept': (self.k_endog, 1),
            'obs_cov': (self.k_endog, self.k_endog, 1),
            'transition': (self.k_states, self.k_states, 1),
            'state_intercept': (self.k_states, 1),
            'selection': (self.k_states, self.k_posdef, 1),
            'state_cov': (self.k_posdef, self.k_posdef, 1),
        }

        # Representation matrices
        # These matrices are only used in the Python object as containers,
        # which will be copied to the appropriate _statespace object if a
        # filter is called.
        scope = locals()
        for name, shape in self.shapes.items():
            if name == 'obs':
                continue
            # Create the initial storage array for each matrix
            setattr(self, '_' + name, np.zeros(shape, dtype=dtype, order="F"))

            # If we were given an initial value for the matrix, set it
            # (notice it is being set via the descriptor)
            if scope[name] is not None:
                setattr(self, name, scope[name])

        # Options
        self.initial_variance = initial_variance
        self.prefix_statespace_map = (statespace_classes
                                      if statespace_classes is not None
                                      else tools.prefix_statespace_map.copy())

        # State-space initialization data
        self.initialization = kwargs.pop('initialization', None)
        basic_inits = ['diffuse', 'approximate_diffuse', 'stationary']

        if self.initialization in basic_inits:
            self.initialize(self.initialization)
        elif self.initialization == 'known':
            if 'constant' in kwargs:
                constant = kwargs.pop('constant')
            elif 'initial_state' in kwargs:
                # TODO deprecation warning
                constant = kwargs.pop('initial_state')
            else:
                raise ValueError('Initial state must be provided when "known"'
                                 ' is the specified initialization method.')
            if 'stationary_cov' in kwargs:
                stationary_cov = kwargs.pop('stationary_cov')
            elif 'initial_state_cov' in kwargs:
                # TODO deprecation warning
                stationary_cov = kwargs.pop('initial_state_cov')
            else:
                raise ValueError('Initial state covariance matrix must be'
                                 ' provided when "known" is the specified'
                                 ' initialization method.')
            self.initialize('known', constant=constant,
                            stationary_cov=stationary_cov)
        elif (not isinstance(self.initialization, Initialization) and
                self.initialization is not None):
            raise ValueError("Invalid state space initialization method.")

        # Check for unused kwargs
        if len(kwargs):
            # raise TypeError(f'{__class__} constructor got unexpected keyword'
            #                 f' argument(s): {kwargs}.')
            msg = (f'Unknown keyword arguments: {kwargs.keys()}.'
                   'Passing unknown keyword arguments will raise a TypeError'
                   ' beginning in version 0.15.')
            warnings.warn(msg, FutureWarning)

        # Matrix representations storage
        self._representations = {}

        # Setup the underlying statespace object storage
        self._statespaces = {}

        # Caches
        self._time_invariant = None

    def __getitem__(self, key):
        _type = type(key)
        # If only a string is given then we must be getting an entire matrix
        if _type is str:
            if key not in self.shapes:
                raise IndexError('"%s" is an invalid state space matrix name'
                                 % key)
            matrix = getattr(self, '_' + key)

            # See note on time-varying arrays, below
            if matrix.shape[-1] == 1:
                return matrix[(slice(None),)*(matrix.ndim-1) + (0,)]
            else:
                return matrix
        # Otherwise if we have a tuple, we want a slice of a matrix
        elif _type is tuple:
            name, slice_ = key[0], key[1:]
            if name not in self.shapes:
                raise IndexError('"%s" is an invalid state space matrix name'
                                 % name)

            matrix = getattr(self, '_' + name)

            # Since the model can support time-varying arrays, but often we
            # will instead have time-invariant arrays, we want to allow setting
            # a matrix slice like mod['transition',0,:] even though technically
            # it should be mod['transition',0,:,0]. Thus if the array in
            # question is time-invariant but the last slice was excluded,
            # add it in as a zero.
            if matrix.shape[-1] == 1 and len(slice_) <= matrix.ndim-1:
                slice_ = slice_ + (0,)

            return matrix[slice_]
        # Otherwise, we have only a single slice index, but it is not a string
        else:
            raise IndexError('First index must the name of a valid state space'
                             ' matrix.')

    def __setitem__(self, key, value):
        _type = type(key)
        # If only a string is given then we must be setting an entire matrix
        if _type is str:
            if key not in self.shapes:
                raise IndexError('"%s" is an invalid state space matrix name'
                                 % key)
            setattr(self, key, value)
        # If it's a tuple (with a string as the first element) then we must be
        # setting a slice of a matrix
        elif _type is tuple:
            name, slice_ = key[0], key[1:]
            if name not in self.shapes:
                raise IndexError('"%s" is an invalid state space matrix name'
                                 % key[0])

            # Change the dtype of the corresponding matrix
            dtype = np.array(value).dtype
            matrix = getattr(self, '_' + name)
            valid_types = ['f', 'd', 'F', 'D']
            if not matrix.dtype == dtype and dtype.char in valid_types:
                matrix = getattr(self, '_' + name).real.astype(dtype)

            # Since the model can support time-varying arrays, but often we
            # will instead have time-invariant arrays, we want to allow setting
            # a matrix slice like mod['transition',0,:] even though technically
            # it should be mod['transition',0,:,0]. Thus if the array in
            # question is time-invariant but the last slice was excluded,
            # add it in as a zero.
            if matrix.shape[-1] == 1 and len(slice_) == matrix.ndim-1:
                slice_ = slice_ + (0,)

            # Set the new value
            matrix[slice_] = value
            setattr(self, name, matrix)
        # Otherwise we got a single non-string key, (e.g. mod[:]), which is
        # invalid
        else:
            raise IndexError('First index must the name of a valid state space'
                             ' matrix.')

    def _clone_kwargs(self, endog, **kwargs):
        """
        Construct keyword arguments for cloning a state space model

        Parameters
        ----------
        endog : array_like
            An observed time-series process :math:`y`.
        **kwargs
            Keyword arguments to pass to the new state space representation
            model constructor. Those that are not specified are copied from
            the specification of the current state space model.
        """

        # We always need the base dimensions, but they cannot change from
        # the base model when cloning (the idea is: if these need to change,
        # need to make a new instance manually, since it's not really cloning).
        kwargs['nobs'] = len(endog)
        kwargs['k_endog'] = self.k_endog
        for key in ['k_states', 'k_posdef']:
            val = getattr(self, key)
            if key not in kwargs or kwargs[key] is None:
                kwargs[key] = val
            if kwargs[key] != val:
                raise ValueError('Cannot change the dimension of %s when'
                                 ' cloning.' % key)

        # Get defaults for time-invariant system matrices, if not otherwise
        # provided
        # Time-varying matrices must be replaced.
        for name in self.shapes.keys():
            if name == 'obs':
                continue

            if name not in kwargs:
                mat = getattr(self, name)
                if mat.shape[-1] != 1:
                    raise ValueError('The `%s` matrix is time-varying. Cloning'
                                     ' this model requires specifying an'
                                     ' updated matrix.' % name)
                kwargs[name] = mat

        # Default is to use the same initialization
        kwargs.setdefault('initialization', self.initialization)

        return kwargs

    def clone(self, endog, **kwargs):
        """
        Clone a state space representation while overriding some elements

        Parameters
        ----------
        endog : array_like
            An observed time-series process :math:`y`.
        **kwargs
            Keyword arguments to pass to the new state space representation
            model constructor. Those that are not specified are copied from
            the specification of the current state space model.

        Returns
        -------
        Representation

        Notes
        -----
        If some system matrices are time-varying, then new time-varying
        matrices *must* be provided.
        """
        kwargs = self._clone_kwargs(endog, **kwargs)
        mod = self.__class__(**kwargs)
        mod.bind(endog)
        return mod

    def extend(self, endog, start=None, end=None, **kwargs):
        """
        Extend the current state space model, or a specific (time) subset

        Parameters
        ----------
        endog : array_like
            An observed time-series process :math:`y`.
        start : int, optional
            The first period of a time-varying state space model to include in
            the new model. Has no effect if the state space model is
            time-invariant. Default is the initial period.
        end : int, optional
            The last period of a time-varying state space model to include in
            the new model. Has no effect if the state space model is
            time-invariant. Default is the final period.
        **kwargs
            Keyword arguments to pass to the new state space representation
            model constructor. Those that are not specified are copied from
            the specification of the current state space model.

        Returns
        -------
        Representation

        Notes
        -----
        This method does not allow replacing a time-varying system matrix with
        a time-invariant one (or vice-versa). If that is required, use `clone`.
        """
        endog = np.atleast_1d(endog)
        if endog.ndim == 1:
            endog = endog[:, np.newaxis]
        nobs = len(endog)

        if start is None:
            start = 0
        if end is None:
            end = self.nobs

        if start < 0:
            start = self.nobs + start
        if end < 0:
            end = self.nobs + end
        if start > self.nobs:
            raise ValueError('The `start` argument of the extension within the'
                             ' base model cannot be after the end of the'
                             ' base model.')
        if end > self.nobs:
            raise ValueError('The `end` argument of the extension within the'
                             ' base model cannot be after the end of the'
                             ' base model.')
        if start > end:
            raise ValueError('The `start` argument of the extension within the'
                             ' base model cannot be after the `end` argument.')

        # Note: if start == end or if end < self.nobs, then we're just cloning
        # (no extension)
        endog = tools.concat([self.endog[:, start:end].T, endog])

        # Extend any time-varying arrays
        error_ti = ('Model has time-invariant %s matrix, so cannot provide'
                    ' an extended matrix.')
        error_tv = ('Model has time-varying %s matrix, so an updated'
                    ' time-varying matrix for the extension period'
                    ' is required.')
        for name, shape in self.shapes.items():
            if name == 'obs':
                continue

            mat = getattr(self, name)

            # If we were *not* given an extended value for this matrix...
            if name not in kwargs:
                # If this is a time-varying matrix in the existing model
                if mat.shape[-1] > 1:
                    # If we have an extension period, then raise an error
                    # because we should have been given an extended value
                    if end + nobs > self.nobs:
                        raise ValueError(error_tv % name)
                    # If we do not have an extension period, then set the new
                    # time-varying matrix to be the portion of the existing
                    # time-varying matrix that corresponds to the period of
                    # interest
                    else:
                        kwargs[name] = mat[..., start:end + nobs]
            elif nobs == 0:
                raise ValueError('Extension is being performed within-sample'
                                 ' so cannot provide an extended matrix')
            # If we were given an extended value for this matrix
            else:
                # TODO: Need to add a check for ndim, and if the matrix has
                # one fewer dimensions than the existing matrix, add a new axis

                # If this is a time-invariant matrix in the existing model,
                # raise an error
                if mat.shape[-1] == 1 and self.nobs > 1:
                    raise ValueError(error_ti % name)

                # Otherwise, validate the shape of the given extended value
                # Note: we do not validate the number of observations here
                # (so we pass in updated_mat.shape[-1] as the nobs argument
                # in the validate_* calls); instead, we check below that we
                # at least `nobs` values were passed in and then only take the
                # first of them as required. This can be useful when e.g. the
                # end user knows the extension values up to some maximum
                # endpoint, but does not know what the calling methods may
                # specifically require.
                updated_mat = np.asarray(kwargs[name])
                if len(shape) == 2:
                    validate_vector_shape(name, updated_mat.shape, shape[0],
                                          updated_mat.shape[-1])
                else:
                    validate_matrix_shape(name, updated_mat.shape, shape[0],
                                          shape[1], updated_mat.shape[-1])

                if updated_mat.shape[-1] < nobs:
                    raise ValueError(error_tv % name)
                else:
                    updated_mat = updated_mat[..., :nobs]

                # Concatenate to get the new time-varying matrix
                kwargs[name] = np.c_[mat[..., start:end], updated_mat]

        return self.clone(endog, **kwargs)

    def diff_endog(self, new_endog, tolerance=1e-10):
        # TODO: move this function to tools?
        endog = self.endog.T
        if len(new_endog) < len(endog):
            raise ValueError('Given data (length %d) is too short to diff'
                             ' against model data (length %d).'
                             % (len(new_endog), len(endog)))
        if len(new_endog) > len(endog):
            nobs_append = len(new_endog) - len(endog)
            endog = np.c_[endog.T, new_endog[-nobs_append:].T * np.nan].T

        new_nan = np.isnan(new_endog)
        existing_nan = np.isnan(endog)
        diff = np.abs(new_endog - endog)
        diff[new_nan ^ existing_nan] = np.inf
        diff[new_nan & existing_nan] = 0.

        is_revision = (diff > tolerance)
        is_new = existing_nan & ~new_nan
        is_revision[is_new] = False

        revision_ix = list(zip(*np.where(is_revision)))
        new_ix = list(zip(*np.where(is_new)))

        return revision_ix, new_ix

    @property
    def prefix(self):
        """
        (str) BLAS prefix of currently active representation matrices
        """
        arrays = (
            self._design, self._obs_intercept, self._obs_cov,
            self._transition, self._state_intercept, self._selection,
            self._state_cov
        )
        if self.endog is not None:
            arrays = (self.endog,) + arrays
        return find_best_blas_type(arrays)[0]

    @property
    def dtype(self):
        """
        (dtype) Datatype of currently active representation matrices
        """
        return tools.prefix_dtype_map[self.prefix]

    @property
    def time_invariant(self):
        """
        (bool) Whether or not currently active representation matrices are
        time-invariant
        """
        if self._time_invariant is None:
            return (
                self._design.shape[2] == self._obs_intercept.shape[1] ==
                self._obs_cov.shape[2] == self._transition.shape[2] ==
                self._state_intercept.shape[1] == self._selection.shape[2] ==
                self._state_cov.shape[2]
            )
        else:
            return self._time_invariant

    @property
    def _statespace(self):
        prefix = self.prefix
        if prefix in self._statespaces:
            return self._statespaces[prefix]
        return None

    @property
    def obs(self):
        r"""
        (array) Observation vector: :math:`y~(k\_endog \times nobs)`
        """
        return self.endog

    def bind(self, endog):
        """
        Bind data to the statespace representation

        Parameters
        ----------
        endog : ndarray
            Endogenous data to bind to the model. Must be column-ordered
            ndarray with shape (`k_endog`, `nobs`) or row-ordered ndarray with
            shape (`nobs`, `k_endog`).

        Notes
        -----
        The strict requirements arise because the underlying statespace and
        Kalman filtering classes require Fortran-ordered arrays in the wide
        format (shaped (`k_endog`, `nobs`)), and this structure is setup to
        prevent copying arrays in memory.

        By default, numpy arrays are row (C)-ordered and most time series are
        represented in the long format (with time on the 0-th axis). In this
        case, no copying or re-ordering needs to be performed, instead the
        array can simply be transposed to get it in the right order and shape.

        Although this class (Representation) has stringent `bind` requirements,
        it is assumed that it will rarely be used directly.
        """
        if not isinstance(endog, np.ndarray):
            raise ValueError("Invalid endogenous array; must be an ndarray.")

        # Make sure we have a 2-dimensional array
        # Note: reshaping a 1-dim array into a 2-dim array by changing the
        #       shape tuple always results in a row (C)-ordered array, so it
        #       must be shaped (nobs, k_endog)
        if endog.ndim == 1:
            # In the case of nobs x 0 arrays
            if self.k_endog == 1:
                endog.shape = (endog.shape[0], 1)
            # In the case of k_endog x 0 arrays
            else:
                endog.shape = (1, endog.shape[0])
        if not endog.ndim == 2:
            raise ValueError('Invalid endogenous array provided; must be'
                             ' 2-dimensional.')

        # Check for valid column-ordered arrays
        if endog.flags['F_CONTIGUOUS'] and endog.shape[0] == self.k_endog:
            pass
        # Check for valid row-ordered arrays, and transpose them to be the
        # correct column-ordered array
        elif endog.flags['C_CONTIGUOUS'] and endog.shape[1] == self.k_endog:
            endog = endog.T
        # Invalid column-ordered arrays
        elif endog.flags['F_CONTIGUOUS']:
            raise ValueError('Invalid endogenous array; column-ordered'
                             ' arrays must have first axis shape of'
                             ' `k_endog`.')
        # Invalid row-ordered arrays
        elif endog.flags['C_CONTIGUOUS']:
            raise ValueError('Invalid endogenous array; row-ordered'
                             ' arrays must have last axis shape of'
                             ' `k_endog`.')
        # Non-contiguous arrays
        else:
            raise ValueError('Invalid endogenous array; must be ordered in'
                             ' contiguous memory.')

        # We may still have a non-fortran contiguous array, so double-check
        if not endog.flags['F_CONTIGUOUS']:
            endog = np.asfortranarray(endog)

        # Set a flag for complex data
        self._complex_endog = np.iscomplexobj(endog)

        # Set the data
        self.endog = endog
        self.nobs = self.endog.shape[1]

        # Reset shapes
        if hasattr(self, 'shapes'):
            self.shapes['obs'] = self.endog.shape

    def initialize(self, initialization, approximate_diffuse_variance=None,
                   constant=None, stationary_cov=None, a=None, Pstar=None,
                   Pinf=None, A=None, R0=None, Q0=None):
        """Create an Initialization object if necessary"""
        if initialization == 'known':
            initialization = Initialization(self.k_states, 'known',
                                            constant=constant,
                                            stationary_cov=stationary_cov)
        elif initialization == 'components':
            initialization = Initialization.from_components(
                a=a, Pstar=Pstar, Pinf=Pinf, A=A, R0=R0, Q0=Q0)
        elif initialization == 'approximate_diffuse':
            if approximate_diffuse_variance is None:
                approximate_diffuse_variance = self.initial_variance
            initialization = Initialization(
                self.k_states, 'approximate_diffuse',
                approximate_diffuse_variance=approximate_diffuse_variance)
        elif initialization == 'stationary':
            initialization = Initialization(self.k_states, 'stationary')
        elif initialization == 'diffuse':
            initialization = Initialization(self.k_states, 'diffuse')

        # We must have an initialization object at this point
        if not isinstance(initialization, Initialization):
            raise ValueError("Invalid state space initialization method.")

        self.initialization = initialization

    def initialize_known(self, constant, stationary_cov):
        """
        Initialize the statespace model with known distribution for initial
        state.

        These values are assumed to be known with certainty or else
        filled with parameters during, for example, maximum likelihood
        estimation.

        Parameters
        ----------
        constant : array_like
            Known mean of the initial state vector.
        stationary_cov : array_like
            Known covariance matrix of the initial state vector.
        """
        constant = np.asarray(constant, order="F")
        stationary_cov = np.asarray(stationary_cov, order="F")

        if not constant.shape == (self.k_states,):
            raise ValueError('Invalid dimensions for constant state vector.'
                             ' Requires shape (%d,), got %s' %
                             (self.k_states, str(constant.shape)))
        if not stationary_cov.shape == (self.k_states, self.k_states):
            raise ValueError('Invalid dimensions for stationary covariance'
                             ' matrix. Requires shape (%d,%d), got %s' %
                             (self.k_states, self.k_states,
                              str(stationary_cov.shape)))

        self.initialize('known', constant=constant,
                        stationary_cov=stationary_cov)

    def initialize_approximate_diffuse(self, variance=None):
        """
        Initialize the statespace model with approximate diffuse values.

        Rather than following the exact diffuse treatment (which is developed
        for the case that the variance becomes infinitely large), this assigns
        an arbitrary large number for the variance.

        Parameters
        ----------
        variance : float, optional
            The variance for approximating diffuse initial conditions. Default
            is 1e6.
        """
        if variance is None:
            variance = self.initial_variance

        self.initialize('approximate_diffuse',
                        approximate_diffuse_variance=variance)

    def initialize_components(self, a=None, Pstar=None, Pinf=None, A=None,
                              R0=None, Q0=None):
        """
        Initialize the statespace model with component matrices

        Parameters
        ----------
        a : array_like, optional
            Vector of constant values describing the mean of the stationary
            component of the initial state.
        Pstar : array_like, optional
            Stationary component of the initial state covariance matrix. If
            given, should be a matrix shaped `k_states x k_states`. The
            submatrix associated with the diffuse states should contain zeros.
            Note that by definition, `Pstar = R0 @ Q0 @ R0.T`, so either
            `R0,Q0` or `Pstar` may be given, but not both.
        Pinf : array_like, optional
            Diffuse component of the initial state covariance matrix. If given,
            should be a matrix shaped `k_states x k_states` with ones in the
            diagonal positions corresponding to states with diffuse
            initialization and zeros otherwise. Note that by definition,
            `Pinf = A @ A.T`, so either `A` or `Pinf` may be given, but not
            both.
        A : array_like, optional
            Diffuse selection matrix, used in the definition of the diffuse
            initial state covariance matrix. If given, should be a
            `k_states x k_diffuse_states` matrix that contains the subset of
            the columns of the identity matrix that correspond to states with
            diffuse initialization. Note that by definition, `Pinf = A @ A.T`,
            so either `A` or `Pinf` may be given, but not both.
        R0 : array_like, optional
            Stationary selection matrix, used in the definition of the
            stationary initial state covariance matrix. If given, should be a
            `k_states x k_nondiffuse_states` matrix that contains the subset of
            the columns of the identity matrix that correspond to states with a
            non-diffuse initialization. Note that by definition,
            `Pstar = R0 @ Q0 @ R0.T`, so either `R0,Q0` or `Pstar` may be
            given, but not both.
        Q0 : array_like, optional
            Covariance matrix associated with stationary initial states. If
            given, should be a matrix shaped
            `k_nondiffuse_states x k_nondiffuse_states`.
            Note that by definition, `Pstar = R0 @ Q0 @ R0.T`, so either
            `R0,Q0` or `Pstar` may be given, but not both.

        Notes
        -----
        The matrices `a, Pstar, Pinf, A, R0, Q0` and the process for
        initializing the state space model is as given in Chapter 5 of [1]_.
        For the definitions of these matrices, see equation (5.2) and the
        subsequent discussion there.

        References
        ----------
        .. [1] Durbin, James, and Siem Jan Koopman. 2012.
           Time Series Analysis by State Space Methods: Second Edition.
           Oxford University Press.
        """
        self.initialize('components', a=a, Pstar=Pstar, Pinf=Pinf, A=A, R0=R0,
                        Q0=Q0)

    def initialize_stationary(self):
        """
        Initialize the statespace model as stationary.
        """
        self.initialize('stationary')

    def initialize_diffuse(self):
        """
        Initialize the statespace model as diffuse.
        """
        self.initialize('diffuse')

    def _initialize_representation(self, prefix=None):
        if prefix is None:
            prefix = self.prefix
        dtype = tools.prefix_dtype_map[prefix]

        # If the dtype-specific representation matrices do not exist, create
        # them
        if prefix not in self._representations:
            # Copy the statespace representation matrices
            self._representations[prefix] = {}
            for matrix in self.shapes.keys():
                if matrix == 'obs':
                    self._representations[prefix][matrix] = (
                        self.obs.astype(dtype)
                    )
                else:
                    # Note: this always makes a copy
                    self._representations[prefix][matrix] = (
                        getattr(self, '_' + matrix).astype(dtype)
                    )
        # If they do exist, update them
        else:
            for matrix in self.shapes.keys():
                existing = self._representations[prefix][matrix]
                if matrix == 'obs':
                    # existing[:] = self.obs.astype(dtype)
                    pass
                else:
                    new = getattr(self, '_' + matrix).astype(dtype)
                    if existing.shape == new.shape:
                        existing[:] = new[:]
                    else:
                        self._representations[prefix][matrix] = new

        # Determine if we need to (re-)create the _statespace models
        # (if time-varying matrices changed)
        if prefix in self._statespaces:
            ss = self._statespaces[prefix]
            create = (
                not ss.obs.shape[1] == self.endog.shape[1] or
                not ss.design.shape[2] == self.design.shape[2] or
                not ss.obs_intercept.shape[1] == self.obs_intercept.shape[1] or
                not ss.obs_cov.shape[2] == self.obs_cov.shape[2] or
                not ss.transition.shape[2] == self.transition.shape[2] or
                not (ss.state_intercept.shape[1] ==
                     self.state_intercept.shape[1]) or
                not ss.selection.shape[2] == self.selection.shape[2] or
                not ss.state_cov.shape[2] == self.state_cov.shape[2]
            )
        else:
            create = True

        # (re-)create if necessary
        if create:
            if prefix in self._statespaces:
                del self._statespaces[prefix]

            # Setup the base statespace object
            cls = self.prefix_statespace_map[prefix]
            self._statespaces[prefix] = cls(
                self._representations[prefix]['obs'],
                self._representations[prefix]['design'],
                self._representations[prefix]['obs_intercept'],
                self._representations[prefix]['obs_cov'],
                self._representations[prefix]['transition'],
                self._representations[prefix]['state_intercept'],
                self._representations[prefix]['selection'],
                self._representations[prefix]['state_cov']
            )

        return prefix, dtype, create

    def _initialize_state(self, prefix=None, complex_step=False):
        # TODO once the transition to using the Initialization objects is
        # complete, this should be moved entirely to the _{{prefix}}Statespace
        # object.
        if prefix is None:
            prefix = self.prefix

        # (Re-)initialize the statespace model
        if isinstance(self.initialization, Initialization):
            if not self.initialization.initialized:
                raise RuntimeError('Initialization is incomplete.')
            self._statespaces[prefix].initialize(self.initialization,
                                                 complex_step=complex_step)
        else:
            raise RuntimeError('Statespace model not initialized.')

class FrozenRepresentation:
    """
    Frozen Statespace Model

    Takes a snapshot of a Statespace model.

    Parameters
    ----------
    model : Representation
        A Statespace representation

    Attributes
    ----------
    nobs : int
        Number of observations.
    k_endog : int
        The dimension of the observation series.
    k_states : int
        The dimension of the unobserved state process.
    k_posdef : int
        The dimension of a guaranteed positive definite
        covariance matrix describing the shocks in the
        measurement equation.
    dtype : dtype
        Datatype of representation matrices
    prefix : str
        BLAS prefix of representation matrices
    shapes : dictionary of name:tuple
        A dictionary recording the shapes of each of
        the representation matrices as tuples.
    endog : ndarray
        The observation vector.
    design : ndarray
        The design matrix, :math:`Z`.
    obs_intercept : ndarray
        The intercept for the observation equation, :math:`d`.
    obs_cov : ndarray
        The covariance matrix for the observation equation :math:`H`.
    transition : ndarray
        The transition matrix, :math:`T`.
    state_intercept : ndarray
        The intercept for the transition equation, :math:`c`.
    selection : ndarray
        The selection matrix, :math:`R`.
    state_cov : ndarray
        The covariance matrix for the state equation :math:`Q`.
    missing : array of bool
        An array of the same size as `endog`, filled
        with boolean values that are True if the
        corresponding entry in `endog` is NaN and False
        otherwise.
    nmissing : array of int
        An array of size `nobs`, where the ith entry
        is the number (between 0 and `k_endog`) of NaNs in
        the ith row of the `endog` array.
    time_invariant : bool
        Whether or not the representation matrices are time-invariant
    initialization : Initialization object
        Kalman filter initialization method.
    initial_state : array_like
        The state vector used to initialize the Kalamn filter.
    initial_state_cov : array_like
        The state covariance matrix used to initialize the Kalamn filter.
    """
    _model_attributes = [
        'model', 'prefix', 'dtype', 'nobs', 'k_endog', 'k_states',
        'k_posdef', 'time_invariant', 'endog', 'design', 'obs_intercept',
        'obs_cov', 'transition', 'state_intercept', 'selection',
        'state_cov', 'missing', 'nmissing', 'shapes', 'initialization',
        'initial_state', 'initial_state_cov', 'initial_variance'
    ]
    _attributes = _model_attributes

    def __init__(self, model):
        # Initialize all attributes to None
        for name in self._attributes:
            setattr(self, name, None)

        # Update the representation attributes
        self.update_representation(model)

    def update_representation(self, model):
        """Update model Representation"""
        # Model
        self.model = model

        # Data type
        self.prefix = model.prefix
        self.dtype = model.dtype

        # Copy the model dimensions
        self.nobs = model.nobs
        self.k_endog = model.k_endog
        self.k_states = model.k_states
        self.k_posdef = model.k_posdef
        self.time_invariant = model.time_invariant

        # Save the state space representation at the time
        self.endog = model.endog
        self.design = model._design.copy()
        self.obs_intercept = model._obs_intercept.copy()
        self.obs_cov = model._obs_cov.copy()
        self.transition = model._transition.copy()
        self.state_intercept = model._state_intercept.copy()
        self.selection = model._selection.copy()
        self.state_cov = model._state_cov.copy()

        self.missing = np.array(model._statespaces[self.prefix].missing,
                                copy=True)
        self.nmissing = np.array(model._statespaces[self.prefix].nmissing,
                                 copy=True)

        # Save the final shapes of the matrices
        self.shapes = dict(model.shapes)
        for name in self.shapes.keys():
            if name == 'obs':
                continue
            self.shapes[name] = getattr(self, name).shape
        self.shapes['obs'] = self.endog.shape

        # Save the state space initialization
        self.initialization = model.initialization

        if model.initialization is not None:
            model._initialize_state()
            self.initial_state = np.array(
                model._statespaces[self.prefix].initial_state, copy=True)
            self.initial_state_cov = np.array(
                model._statespaces[self.prefix].initial_state_cov, copy=True)
            self.initial_diffuse_state_cov = np.array(
                model._statespaces[self.prefix].initial_diffuse_state_cov,
                copy=True)


import torch
import torch.nn.functional as F
from tqdm import tqdm
import numpy as np

class LSTMModel(torch.nn.Module):
    def __init__(self, diag_vocab_size, med_vocab_size, diag_embedding_size,
                 med_embedding_size, diag_hidden_size, med_hidden_size, hidden_size,
                 bidirectional=True):  # end_index, pad_index,
        super().__init__()
        # self.pad_index = pad_index
        # self.end_index = end_index
        self.diag_embedding = torch.nn.Linear(diag_vocab_size, diag_embedding_size, bias=False)
        self.med_embedding = torch.nn.Linear(med_vocab_size, med_embedding_size, bias=False)

        self.diag_encoder = torch.nn.LSTM(diag_embedding_size, diag_hidden_size, batch_first=True,
                                          bidirectional=bidirectional)  # , dropout=0.5)  # 1 layer, dropout not used

        self.med_encoder = torch.nn.LSTM(med_embedding_size, med_hidden_size, batch_first=True,
                                         bidirectional=bidirectional)  # , dropout=0.5)

        if bidirectional:
            diag_hidden_size = diag_hidden_size * 2
            med_hidden_size = med_hidden_size * 2

        self.attention_diag_encoder = torch.nn.Sequential(
            torch.nn.Linear(diag_hidden_size, 1),
            torch.nn.Tanh(),
        )
        self.attention_med_encoder = torch.nn.Sequential(
            torch.nn.Linear(med_hidden_size, 1),
            torch.nn.Tanh(),
        )

        # IPW   + 2 --> 2 dims of age and sex of demo features
        self.lstm2hidden_ipw = torch.nn.Sequential(
            torch.nn.Linear(med_hidden_size + diag_hidden_size + 3, hidden_size),
            torch.nn.ReLU(),
        )
        self.hidden2out_ipw = torch.nn.Linear(hidden_size, 1,
                                              bias=False)  # outputs logits, not using sigmoid here for probability

        # Outcome
        # self.lstm2hidden_outcome = torch.nn.Sequential(
        #     torch.nn.Linear(med_hidden_size + diag_hidden_size + 3, hidden_size),
        #     torch.nn.ReLU(),
        # )
        # self.hidden2out_outcome = torch.nn.Linear(hidden_size, 1, bias=False)

    def softmax_masked(self, inputs, mask, dim=1, epsilon=0.0000001):
        inputs_exp = torch.exp(inputs)
        inputs_exp = inputs_exp * mask.float()
        inputs_exp_sum = inputs_exp.sum(dim=dim)
        inputs_attention = inputs_exp / (inputs_exp_sum.unsqueeze(dim) + epsilon)

        return inputs_attention

    def diag_encode(self, inputs):
        inputs_mask = (inputs.sum(dim=-1) != 0).long()
        inputs_emb = self.diag_embedding(inputs.float())  # change this to graph convolution later
        outputs, (h, c) = self.diag_encoder(inputs_emb)

        att_enc = self.attention_diag_encoder(outputs).squeeze(-1)
        att_normalized = self.softmax_masked(att_enc, inputs_mask)
        hidden = torch.sum(outputs * att_normalized.unsqueeze(-1), dim=1)
        original = torch.sum(inputs.float() * att_normalized.unsqueeze(-1), dim=1)

        return hidden, original

    def med_encode(self, inputs):  # bs * Time * Dim
        inputs_mask = (inputs.sum(dim=-1) != 0).long()  # bs * Time  masking what?
        inputs_emb = self.med_embedding(
            inputs.float())  # bs * Time * HiddenDim  # change this to graph convolution later
        outputs, (h, c) = self.med_encoder(inputs_emb)  # Outputs: bs * Time * HiddenDim

        att_enc = self.attention_med_encoder(outputs).squeeze(-1)  # bs * Time
        att_normalized = self.softmax_masked(att_enc, inputs_mask)  # bs * Time
        hidden = torch.sum(outputs * att_normalized.unsqueeze(-1), dim=1)  # bs * HiddenDim
        original = torch.sum(inputs.float() * att_normalized.unsqueeze(-1),
                             dim=1)  # (bs * Time * Dim) * (bs * Time) attention over time --> (bs * Dim) time attention weighted sum of baseline features

        return hidden, original

    def forward(self, confounder):
        diag_inputs, med_inputs, sexes, ages, days = confounder
        diag_hidden, diag_original = self.diag_encode(diag_inputs)  # bs * T1 * D1 --> bs * hidden1, bs * D1
        med_hidden, med_original = self.med_encode(med_inputs)  # bs * T2 * D2 --> bs * hidden2, bs * D2
        original = torch.cat((diag_original, med_original,
                              sexes.float().view(sexes.size(0), 1),
                              ages.float().view(ages.size(0), 1),
                              days.float().view(days.size(0), 1)), dim=1)  # bs * (D1+D2+2)

        # IPW
        hidden = torch.cat((diag_hidden, med_hidden,
                            sexes.float().view(sexes.size(0), 1),
                            ages.float().view(ages.size(0), 1),
                            days.float().view(days.size(0), 1)), dim=1)  # bs * (hidden1+hidden2+2)
        hidden = self.lstm2hidden_ipw(hidden)  # bs * (hidden1+hidden2+2) --> bs * hidden3
        outputs_logits_ipw = self.hidden2out_ipw(hidden)  # bs * hidden3 --> bs * 1

        return outputs_logits_ipw.view(outputs_logits_ipw.size(0)), original

def lstm_fit(model, optimizer, epochs, train_loader_shuffled, cuda):
    for epoch in tqdm(range(epochs)):
        epoch_losses_ipw = []
        for confounder, treatment, outcome in train_loader_shuffled:
            model.train()
            # train IPW
            optimizer.zero_grad()

            if cuda:  # confounder = (diag, med, sex, age)
                for ii in range(len(confounder)):
                    confounder[ii] = confounder[ii].to('cuda')
                treatment = treatment.to('cuda')

            treatment_logits, _ = model(confounder)
            loss_ipw = F.binary_cross_entropy_with_logits(treatment_logits, treatment.float())

            loss_ipw.backward()
            optimizer.step()
            epoch_losses_ipw.append(loss_ipw.item())

        # just finish 1 epoch
        # scheduler.step()
        epoch_losses_ipw = np.mean(epoch_losses_ipw)

        print('epoch:{}, loss:{}'.format(epoch, epoch_losses_ipw))

    return model, epoch_losses_ipw

# def lstm_cross_validation_fit(dataset, indices, kfold=10, verbose=1, shuffle=True):


"""
Author: Manish Saroya
Contact: saroyam@oregonstate.edu
"""
import networkx as nx
import numpy as np
from persistence.utils import load_hilbert_map, convert_map_dict_to_array, collision_check
import matplotlib.pyplot as pl
from sklearn.neighbors import NearestNeighbors
from bresenham import bresenham
import os
import argparse
import datetime
import pickle
from math import sqrt, log
import time

def save_img(data, graph, dir, save_data=True, save_graph=True, num_sample=0):
	"""
	:param data: map data dictionary
	:param graph: prm nx graph
	:param dir: output directory path
	:return: save prm graph and visualization image to output directory
	"""
	fig = pl.figure(figsize=(40/4, 35/4))
	ax = fig.add_subplot(111)
	pl.axis("equal")
	pl.scatter(data['Xq'][:, 0], data['Xq'][:, 1], c=data['yq'], cmap='viridis', s=70, vmin=0, vmax=1, edgecolors='')
	pl.colorbar()
	position = nx.get_node_attributes(graph, 'pos')
	for node_1, node_2 in graph.edges:
		weights = np.concatenate([[position[node_1]], [position[node_2]]])
		line, = pl.plot(*weights.T, color='lightsteelblue')
		pl.setp(line, linewidth=2, color='lightsteelblue')
	#pl.title('Test image')
	if save_data:
		pl.savefig(dir + str(num_sample) + "prm.png")
	if save_graph:
		with open(dir + str(num_sample) + 'prm.pickle', 'wb') as handle:
			pickle.dump(graph, handle)

def hilbert_samples(map_data, exp_factor, obs_threshold, num_samples=600):
	"""
	:param map_data: map info
	:param exp_factor: exponential factor to sample more from free space
	:param num_samples: number of samples required
	:return: samples
	"""
	map_data['yq'] = np.ones(len(map_data['yq'])) - map_data['yq']
	#map_data['yq'] = np.exp(exp_factor * map_data['yq'])
	map_data["yq"][map_data["yq"] < (1-obs_threshold)] = 0
	# normalize the probabilities
	map_data['yq'] /= np.linalg.norm(map_data['yq'], ord=1)

	samples_list = map_data['Xq'][np.random.choice(len(map_data['Xq']), size=num_samples, p=map_data['yq'])]
	return samples_list

if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument('--number_of_samples', type=int, default=15000)
	parser.add_argument('--exp_factor', type=int, default=20)
	parser.add_argument('--obstacle_threshold', type=float, default=0.25)
	parser.add_argument('--max_nodes', type=int, default=4000)
	parser.add_argument('--k_nearest', type=int, default=3900)
	parser.add_argument('--log_dir', type=str, default='./output')
	parser.add_argument('--connection_radius', type=float, default=5.0)
	parser.add_argument('--map_type', type=str, default="intel")
	parser.add_argument('--gamma_prm', type=float, default=30.0)
	#EDIT PRM STAR PARAM
	args = parser.parse_args()
	args.log_dir = './output/max_nodes-' + args.map_type + str(args.max_nodes) + "-obs-thres" + str(args.obstacle_threshold) +\
				   "-k_nearest-" + \
				   str(args.k_nearest) + "-connection_radius-" + str(args.connection_radius) + "-date-" + \
				   datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") + '/'

	if not os.path.exists(args.log_dir):
		os.makedirs(args.log_dir)
	tic = time.time()
	# load map
	map_data, resolution = load_hilbert_map(map_type=args.map_type)
	#resolution = 0.2
	#with open("freiburg_ground_map_q_resolution_final.pickle", 'rb') as tf:
	#	map_data = pickle.load(tf)
	map_array = convert_map_dict_to_array(map_data, resolution)
	map_data["yq"] = 1.0 * (map_data["yq"] > args.obstacle_threshold)
	# get samples from hilbert maps
	sample_list = hilbert_samples(map_data.copy(), args.exp_factor, args.obstacle_threshold, num_samples=args.number_of_samples)
	# take unique samples
	sample_list = [list(t) for t in set(tuple(element) for element in sample_list)]
	# truncated based on max nodes
	sample_list = sample_list[:args.max_nodes]
	# find k nearest neighbor
	nbrs = NearestNeighbors(n_neighbors=args.k_nearest, algorithm='ball_tree').fit(sample_list)
	distances, indices = nbrs.kneighbors(sample_list)
	# create gragh
	prm_graph = nx.Graph()
	# add graph nodes
	for indx, s in enumerate(sample_list):
		prm_graph.add_node(indx, pos=(s[0], s[1]))
	# add graph edges
	for row, node_adjacency_list in enumerate(indices):
		args.connection_radius = args.gamma_prm * sqrt(log(row + 1) / (row+1))
		#print(row, args.connection_radius)
		for column, other_node in enumerate(node_adjacency_list):
			distance_metric = distances[row][column] < args.connection_radius
			if distance_metric:
				validation_metric = (node_adjacency_list[0] != other_node) and (node_adjacency_list[0] > other_node)
				if validation_metric:
					collision_metric = collision_check(map_array, sample_list[node_adjacency_list[0]],
													   sample_list[other_node], args.obstacle_threshold, resolution)
					if collision_metric:
						prm_graph.add_edge(node_adjacency_list[0], other_node, distance=distances[row][column])
			#if distance_metric and collision_metric and validation_metric:
				#prm_graph.add_edge(node_adjacency_list[0], other_node, distance=distances[row][column])
		if row == 500 or row == 1000 or row == 2000:
			print("rendering", str(row))
			print("generation time ", row , args.map_type, " ", time.time()-tic)
			save_img(map_data, prm_graph, args.log_dir, num_sample=row)
	print("rendering")
	print("generation time ", str(4000), args.map_type, " ", time.time() - tic)
	#save_img(map_data, prm_graph, args.log_dir, num_sample=4000)


# Loading all necessary libraries

import matplotlib.pyplot as plt
import pandas as pd
import re
import csv
from wordcloud import WordCloud, STOPWORDS, ImageColorGenerator
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from collections import Counter
set(stopwords.words('english'))

# We are getting our data for categories and polarity, for the further sentiment analysis
category = pd.read_csv(r"C:/Users/Diana/Desktop/AUA/hpc/category.csv", header=None)
polarity = pd.read_csv(r"C:/Users/Diana/Desktop/AUA/hpc/polarity.csv", header=None)
df = pd.concat([category, polarity], axis=1)
df.columns = ['Category', 'Polarity']

# Here we are plotting the percentage of tweets for each category

fig1 = plt.figure()
df.Category.value_counts().plot(kind='pie', autopct='%1.0f%%')
plt.title('The percentage of tweets for each category')
fig1.savefig(r"C:/Users/Diana/Desktop/AUA/hpc/img1.png")
plt.show()
plt.close(fig1)

# This snippet of code generates the plot for the distribution of sentiments across all the tweets

fig2 = plt.figure()
df.Polarity.value_counts().plot(kind='pie', autopct='%1.0f%%', colors=["red", "yellow", "green"])
plt.title('The distribution of sentiments across all the tweets')
plt.ylabel(' ')
fig2.savefig(r"C:/Users/Diana/Desktop/AUA/hpc/img2.png")
plt.show()

# Finally, we draw the distribution of sentiment for each category
df.groupby(['Category', 'Polarity']).Polarity.count().unstack().plot(kind='bar')
plt.title('The distribution of sentiment for each category')
plt.grid(False)
plt.savefig(r"C:/Users/Diana/Desktop/AUA/hpc/img3.png")
plt.show()

# Cleaning our data from tweets

def clean_tweet(tweet):
    tweet = re.sub('http\S+\s*', '', tweet)  # remove URLs
    tweet = re.sub('RT|cc', '', tweet)  # remove RT and cc
    tweet = re.sub('#\S+', '', tweet)  # remove hash tags
    tweet = re.sub('@\S+', '', tweet)  # remove mentions
    tweet = re.sub(r"\s+[a-zA-Z]\s+", " ", tweet)
    tweet = re.sub('[%s]' % re.escape("""!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~"""), '', tweet)  # remove punctuations
    tweet = re.sub('\s+', ' ', tweet)  # remove extra whitespace
    tweet = re.sub('(?<=[a-z])\'(?=[a-z])', '', tweet)
    tweet = tweet.replace('’', '')  # remove apostrophes
    return tweet

def remove_emoji(string):
    emoji_pattern = re.compile("["
                               u"\U0001F600-\U0001F64F"  # emoticons
                               u"\U0001F300-\U0001F5FF"  # symbols & pictographs
                               u"\U0001F680-\U0001F6FF"  # transport & map symbols
                               u"\U0001F1E0-\U0001F1FF"  # flags (iOS)
                               u"\U00002702-\U000027B0"
                               u"\U000024C2-\U0001F251"

                               "]+", flags=re.UNICODE)

    return emoji_pattern.sub(r'', string)

# We get our data from the csv file then
# We are removing all emoticons and other redundant characters from our data
array = []
with open(r"C:/Users/Diana/Desktop/AUA/hpc/tweets.csv", encoding='utf-8-sig') as csv_file:
    df = csv.reader(csv_file)
    for row in df:
        row[0] = clean_tweet(row[0])
        row[0] = remove_emoji(row[0])
        row[0] = row[0].lower()
        array.append(row[0])

# Removing stopwords from the data set
data_set = ''.join(array)
stop_words = set(stopwords.words('english'))
word_tokens = word_tokenize(data_set)
filtered_data = [w for w in word_tokens if not w in stop_words]
filtered_data = []

for w in word_tokens:
    if w not in stop_words:
        filtered_data.append(w)

final_data = ' '.join(filtered_data)

# Now we are getting 10 most frequent words from our cleaned data
# split() returns list of all the words in the string
split_it = final_data.split()

# We pass the split_it list to instance of Counter class, this counts the occurrence of each word in the string.
Counter = Counter(split_it)

# most_common(10) produces 10 frequently encountered words and their respective counts.
most_occur = Counter.most_common(10)

# We specify labels and frequencies for plotting the graph of 10 frequently encountered words.
label = []
frequency = []
for word in most_occur:
    label.append(word[0].capitalize())
    frequency.append(word[1])

# Here we are plotting the graph
labels = label[0:10]
fig4 = plt.figure(figsize=(10, 5))
plt.bar(labels, frequency[0:10], width=0.8, alpha=0.5, capsize=8, linewidth=1)
plt.xlabel("10 more frequent words", size=14)
plt.ylabel("Frequency", size=14)
plt.title('Word Frequencies', size=18)
plt.grid(False)
plt.gca().spines["top"].set_visible(False)
plt.gca().spines["right"].set_visible(False)
fig4.savefig(r"C:/Users/Diana/Desktop/AUA/hpc/img4.png")
plt.show()
plt.close(fig4)

# Now we create and generate a word cloud image.
plt.figure(figsize=(10, 10))
wordcloud = WordCloud(width=800, height=500, max_font_size=80, max_words=150, background_color="white",
                      contour_color="steelblue", colormap="nipy_spectral").generate(final_data)
plt.imshow(wordcloud, interpolation="hermite")
plt.axis("off")
plt.show()

# Now we save the image
wordcloud.to_file(r"C:/Users/Diana/Desktop/AUA/hpc/wordcloud_img.png")