import sys

import re
import pandas as pd

from datetime import datetime, timedelta
import sqlite3
import time

MAPS_TLAS = { 'Overpass' : 'ovp',
'Cache' : 'cch',
'Nuke' : 'nuke',
'Train' : 'trn',
'Mirage' : 'mrg',
'Cobblestone' : 'cbl',
'Inferno' : 'inf',
'Dust2' : 'd2',
'Tuscan' : 'tcn',
'Season' : 'ssn'
}

tmpidx = ['MatchID',
  'MapName',
  'Event',
  'MatchDate',
  'GOTV_link',
  'TeamLeft',
  'TeamRight',
  'TeamLeftFirstKills',
  'TeamRightFirstKills',
  'TeamLeftClutches',
  'TeamRightClutches',
  'TeamLeftWinPistol',
  'TeamRightWinPistol',
  'TeamLeftTElims',
  'TeamRightTElims',
  'TeamLeftCTElims',
  'TeamRightCTElims',
  'TeamLeftBombDefusals',
  'TeamRightBombDefusals',
  'TeamLeftBombExploded',
  'TeamRightBombExploded',
  'TeamLeftTWins',
  'TeamRightTWins',
  'TeamLeftCTWins',
  'TeamRightCTWins',
  'TeamLeftTotalRounds',
  'TeamRightTotalRounds'
]

def getMonthPerformance(team, match_date):
    md = match_date.isoformat()[:10]
    md_1month = (match_date - timedelta(days = 31)).isoformat()[:10]
    md_3month = (match_date - timedelta(days = 30*3)).isoformat()[:10]
    md_6month = (match_date - timedelta(days = 30*6)).isoformat()[:10]

    maps_1m = c.execute("SELECT * FROM maps WHERE (TeamLeft = ? OR TeamRight = ?) AND MatchDate > ?  AND MatchDate < ?", (team, team, md_1month, md)).fetchall()

    if maps_1m == []:
        TRounds_1m = "NA"
        CTRounds_1m = "NA"
        TotalRounds_1m = "NA"
        Wins_1m = "NA"
        TRoundsRate_1m = "NA"
        CTRoundsRate_1m = "NA"
        TotalRoundsRate_1m = "NA"
        WinsRate_1m = "NA"
    else:
        TRounds_1m = 0
        CTRounds_1m = 0
        TotalRounds_1m = 0
        Wins_1m = 0
        count = 0
        for m in maps_1m:
            count += 1
            if m[idx['TeamLeft']] == team:
                TRounds_1m += m[idx['TeamLeftTWins']]
                CTRounds_1m += m[idx['TeamLeftCTWins']]
                TotalRounds_1m += m[idx['TeamLeftTotalRounds']]
                Wins_1m += (1 if (m[idx['TeamLeftTotalRounds']] >= 16) else 0)
            else:
                TRounds_1m += m[idx['TeamRightTWins']]
                CTRounds_1m += m[idx['TeamRightCTWins']]
                TotalRounds_1m += m[idx['TeamRightTotalRounds']]
                Wins_1m += (1 if (m[idx['TeamRightTotalRounds']] >= 16) else 0)
        TRoundsRate_1m = TRounds_1m / count
        CTRoundsRate_1m = CTRounds_1m / count
        TotalRoundsRate_1m = TotalRounds_1m / count
        WinsRate_1m = Wins_1m / count

    maps_3m = c.execute("SELECT * FROM maps WHERE (TeamLeft = ? OR TeamRight = ?) AND MatchDate > ?  AND MatchDate < ?", (team, team, md_3month, md)).fetchall()

    if maps_3m == []:
        TRounds_3m = "NA"
        CTRounds_3m = "NA"
        TotalRounds_3m = "NA"
        Wins_3m = "NA"
        TRoundsRate_3m = "NA"
        CTRoundsRate_3m = "NA"
        TotalRoundsRate_3m = "NA"
        WinsRate_3m = "NA"
    else:
        TRounds_3m = 0
        CTRounds_3m = 0
        TotalRounds_3m = 0
        Wins_3m = 0
        count = 0
        for m in maps_3m:
            count += 1
            if m[idx['TeamLeft']] == team:
                TRounds_3m += m[idx['TeamLeftTWins']]
                CTRounds_3m += m[idx['TeamLeftCTWins']]
                TotalRounds_3m += m[idx['TeamLeftTotalRounds']]
                Wins_3m += (1 if (m[idx['TeamLeftTotalRounds']] >= 16) else 0)
            else:
                TRounds_3m += m[idx['TeamRightTWins']]
                CTRounds_3m += m[idx['TeamRightCTWins']]
                TotalRounds_3m += m[idx['TeamRightTotalRounds']]
                Wins_3m += (1 if (m[idx['TeamRightTotalRounds']] >= 16) else 0)
        TRoundsRate_3m = TRounds_3m / count
        CTRoundsRate_3m = CTRounds_3m / count
        TotalRoundsRate_3m = TotalRounds_3m / count
        WinsRate_3m = Wins_3m / count

    maps_6m = c.execute("SELECT * FROM maps WHERE (TeamLeft = ? OR TeamRight = ?) AND MatchDate > ?  AND MatchDate < ?", (team, team, md_6month, md)).fetchall()

    if maps_6m == []:
        TRounds_6m = "NA"
        CTRounds_6m = "NA"
        TotalRounds_6m = "NA"
        Wins_6m = "NA"
        TRoundsRate_6m = "NA"
        CTRoundsRate_6m = "NA"
        TotalRoundsRate_6m = "NA"
        WinsRate_6m = "NA"
    else:
        TRounds_6m = 0
        CTRounds_6m = 0
        TotalRounds_6m = 0
        Wins_6m = 0
        count = 0
        for m in maps_6m:
            count += 1
            if m[idx['TeamLeft']] == team:
                TRounds_6m += m[idx['TeamLeftTWins']]
                CTRounds_6m += m[idx['TeamLeftCTWins']]
                TotalRounds_6m += m[idx['TeamLeftTotalRounds']]
                Wins_6m += (1 if (m[idx['TeamLeftTotalRounds']] >= 16) else 0)
            else:
                TRounds_6m += m[idx['TeamRightTWins']]
                CTRounds_6m += m[idx['TeamRightCTWins']]
                TotalRounds_6m += m[idx['TeamRightTotalRounds']]
                Wins_6m += (1 if (m[idx['TeamRightTotalRounds']] >= 16) else 0)
        TRoundsRate_6m = TRounds_6m / count
        CTRoundsRate_6m = CTRounds_6m / count
        TotalRoundsRate_6m = TotalRounds_6m / count
        WinsRate_6m = Wins_6m / count

    return {'TRounds_1m' : TRounds_1m, 'TRoundsRate_1m' : TRoundsRate_1m, \
        'CTRounds_1m' : CTRounds_1m, 'CTRoundsRate_1m' : CTRoundsRate_1m, \
        'TotalRounds_1m': TotalRounds_1m, 'TotalRoundsRate_1m':TotalRoundsRate_1m, \
        'Wins_1m' : Wins_1m, 'WinsRate_1m' : WinsRate_1m, \
        'TRounds_3m' :  TRounds_3m, 'TRoundsRate_3m': TRoundsRate_3m, \
        'CTRounds_3m' : CTRounds_3m, 'CTRoundsRate_3m' : CTRoundsRate_3m,\
        'TotalRounds_3m' : TotalRounds_3m, 'TotalRoundsRate_3m' : TotalRoundsRate_3m,\
        'Wins_3m' : Wins_3m, 'WinsRate_3m' : WinsRate_3m, \
        'TRounds_6m' : TRounds_6m, 'TRoundsRate_6m' : TRoundsRate_6m, \
        'CTRounds_6m' : CTRounds_6m, 'CTRoundsRate_6m' : CTRoundsRate_6m, \
        'TotalRounds_6m' : TotalRounds_6m, 'TotalRoundsRate_6m' : TotalRoundsRate_6m, \
        'Wins_6m' : Wins_6m, 'WinsRate_6m' : WinsRate_6m}

idx = {}

for i in tmpidx:
    idx[i] = tmpidx.index(i)

numRows = 0

dataMatrix = None

# team lifetime winrate on map
conn = sqlite3.connect('csgo.db')
c = conn.cursor()

maps = c.execute("SELECT * FROM maps")

for m in maps.fetchall():
    # map affinity

    print( m[idx['TeamLeftTotalRounds']] , m[idx['TeamRightTotalRounds']] )
    if m[idx['TeamLeftTotalRounds']] > m[idx['TeamRightTotalRounds']]:
        y = 1
    elif m[idx['TeamLeftTotalRounds']] < m[idx['TeamRightTotalRounds']]:
        y = -1
    else:
        y = 0
        assert m[idx['TeamLeftTotalRounds']] == m[idx['TeamRightTotalRounds']]

    mapname = m[idx['MapName']]
    teamLeft = m[idx['TeamLeft']]
    teamRight = m[idx['TeamRight']]
    matchDate = m[idx['MatchDate']]
    matchDateText = matchDate.split()[0]
    matchTime = int(matchDate.split()[1][:2])  #hour digits of date
    matchDate = datetime.strptime(matchDate.split()[0] , "%Y-%m-%d")

    # team rating
    teamLeftRating = c.execute("SELECT * FROM team_history WHERE Team_ID = ?", (teamLeft,))
    closestDate = None
    closestEntry = None
    for entry in teamLeftRating.fetchall():
        d = datetime.strptime(entry[1], "%Y-%m-%d")
        if matchDate - d > timedelta(days=0):
            if closestDate is None:
                closestDate = d
                closestEntry = entry
            elif matchDate - d < matchDate - closestDate:
                closestDate = d
                closestEntry = entry
    if closestDate is None:
        teamLeftRating = 1.0
        teamLeftKD = 1.0
        teamLeftDiff = 0
    else:
        teamLeftRating = closestEntry[5]
        teamLeftKD = closestEntry[4]
        teamLeftDiff = closestEntry[3]

    teamRightRating = c.execute("SELECT * FROM team_history WHERE Team_ID = ?", (teamRight,))
    closestDate = None
    closestEntry = None
    for entry in teamRightRating.fetchall():
        d = datetime.strptime(entry[1], "%Y-%m-%d")
        if matchDate - d > timedelta(days=0):
            if closestDate is None:
                closestDate = d
                closestEntry = entry
            elif matchDate - d < matchDate - closestDate:
                closestDate = d
                closestEntry = entry
    if closestDate is None:
        teamRightRating = 1.0
        teamRightKD = 1.0
        teamRightDiff = 0
    else:
        teamRightRating = closestEntry[5]
        teamRightKD = closestEntry[4]
        teamRightDiff = closestEntry[3]

    # is new team
    newTeam = c.execute("SELECT MatchDate FROM maps WHERE TeamLeft = ? OR TeamRight = ? ORDER BY MatchDate ASC LIMIT 1", (teamLeft, teamLeft))
    x = newTeam.fetchone()
    if x is None:
        newTeamLeft = True
    else:
        if datetime.strptime(x[0].split()[0] , "%Y-%m-%d") - matchDate > timedelta(days=60):
            newTeamLeft = False
        else:
            newTeamLeft = True

    newTeam = c.execute("SELECT MatchDate FROM maps WHERE TeamLeft = ? OR TeamRight = ? ORDER BY MatchDate ASC LIMIT 1", (teamRight, teamRight))
    x = newTeam.fetchone()
    if x is None:
        newTeamRight = True
    else:
        if datetime.strptime(x[0].split()[0] , "%Y-%m-%d") - matchDate > timedelta(days=60):
            newTeamRight = False
        else:
            newTeamRight = True

    # time since last match
    lastMatch = c.execute("SELECT MatchDate FROM maps WHERE MatchDate < ? and ( TeamLeft = ? OR TeamRight = ? ) ORDER BY MatchDate DESC LIMIT 1", (matchDateText, teamLeft, teamLeft))
    x = lastMatch.fetchone()
    if x is None:
        teamLeftLastMatch = "NA"
    else:
        teamLeftLastMatch = (matchDate - datetime.strptime(x[0].split()[0] , "%Y-%m-%d")).days

    lastMatch = c.execute("SELECT MatchDate FROM maps WHERE MatchDate < ? and (TeamLeft = ? OR TeamRight = ?) ORDER BY MatchDate DESC LIMIT 1", (matchDateText, teamRight, teamRight))
    x = lastMatch.fetchone()
    if x is None:
        teamRightLastMatch = "NA"
    else:
        teamRightLastMatch = (matchDate - datetime.strptime(x[0].split()[0] , "%Y-%m-%d")).days

    #matches this week
    matchDateMinusWeek = (matchDate - timedelta(days=7)).isoformat()[:10]
    teamLeftThisWeek = c.execute("SELECT COUNT(*) FROM maps WHERE (TeamLeft = ? OR TeamRight = ?) AND MatchDate > ? and  MatchDate < ?",  (teamLeft, teamLeft, matchDateMinusWeek, matchDate)   ).fetchone()[0]
    teamRightThisWeek = c.execute("SELECT COUNT(*) FROM maps WHERE (TeamLeft = ? OR TeamRight = ?) AND MatchDate > ? and MatchDate < ?",  (teamRight, teamRight, matchDateMinusWeek, matchDate)   ).fetchone()[0]

    # map affinity
    map_affinity_left = c.execute("SELECT Winrate FROM map_affinity WHERE Team_ID = ? AND MapName = ?" , ( teamLeft, mapname)).fetchone()
    if map_affinity_left is None:
        map_affinity_left = .5
    else:
        map_affinity_left = map_affinity_left[0]
    map_affinity_right = c.execute("SELECT Winrate FROM map_affinity WHERE Team_ID = ? AND MapName = ?" , ( teamRight, mapname)).fetchone()
    if map_affinity_right is None:
        map_affinity_right = .5
    else:
        map_affinity_right = map_affinity_right[0]

    #map count
    teamLeftMapCount = int(c.execute("SELECT COUNT(*) FROM maps WHERE TeamLeft = ?", (teamLeft,)).fetchone()[0] ) + \
        int(c.execute("SELECT count(*) FROM maps WHERE TeamRight = ?", (teamLeft,)).fetchone()[0] )

    teamRightMapCount = int(c.execute("SELECT COUNT(*) FROM maps WHERE TeamLeft = ?", (teamRight,)).fetchone()[0] ) + \
        int(c.execute("SELECT COUNT(*) FROM maps WHERE TeamRight = ?", (teamRight,)).fetchone()[0] )

    teamLeftMonthlyPerf = getMonthPerformance(teamLeft, matchDate)
    teamRightMonthlyPerf = getMonthPerformance(teamRight, matchDate)

    prflen = len(teamLeftMonthlyPerf)

#    print(mapname, teamLeft, teamRight, matchDateText, matchTime, teamLeftRating , \
#        teamRightRating, teamLeftKD, teamLeftDiff, newTeamLeft, newTeamRight, \
#        teamLeftLastMatch, teamRightLastMatch, teamLeftThisWeek, map_affinity_left, teamLeftMapCount)

    print(numRows)
    numRows += 1

    if dataMatrix is None:
        columns = ['mapname', 'matchTime', 'teamLeftRating', 'teamRightRating', 'teamLeftKD', 'teamRightKD', \
            'teamLeftDiff', 'teamRightDiff', 'newTeamLeft', 'newTeamRight', \
            'teamLeftLastMatch', 'teamRightLastMatch', 'teamLeftThisWeek', 'teamRightThisWeek', \
            'map_affinity_left', 'map_affinity_right']

        columns += ['teamLeft' + i for i in teamLeftMonthlyPerf]
        columns += ['teamRight' + i for i in teamRightMonthlyPerf]
        columns += ['winner']
        dataMatrix = []

    datarow = [mapname, matchTime,  teamLeftRating ,  teamRightRating ,  teamLeftKD ,  teamRightKD , \
         teamLeftDiff ,  teamRightDiff ,  newTeamLeft ,  newTeamRight , \
         teamLeftLastMatch ,  teamRightLastMatch ,  teamLeftThisWeek ,  teamRightThisWeek , \
         map_affinity_left ,  map_affinity_right]

    datarow += [teamLeftMonthlyPerf[i[8:]] for i in columns[16:16+prflen]]
    datarow += [teamRightMonthlyPerf[i[9:]] for i in columns[16+prflen:-1]]
    datarow += [y]
    dataRow = dict(zip(columns, datarow))
    dataMatrix.append(dataRow)

dataMatrix =pd.DataFrame(dataMatrix)

print(dataMatrix)
import pickle
pickle.dump(dataMatrix, open('dataMatrix', 'wb'), pickle.HIGHEST_PROTOCOL)


import os
import torch
from lib.dataset.landsat_8_dataloader import Get_dataloader
from lib.model.create_generator import create_generator
from lib.model.generator.ms_net import MS_Net
from lib.model.generator.msa_net import MSA_Net
from lib.model.generator.rsc_net import RSC_Net
from lib.model.generator.ssa_net import SSA_Net
from lib.utils.config import Config
from lib.utils.parse_args import parse_args
from rm_cloud.single_model_methods.tester import Testner

def eval():
    config = Config().parse()
    args = parse_args(config)
    print(args)
    train_loader, test_loader = Get_dataloader(args)
    os.environ['CUDA_VISIBLE_DEVICES'] = args.devices

    if args.model.arch == 'rsc':
        model = RSC_Net()
    if args.model.arch == 'msa':
        model = MSA_Net()
    if args.model.arch == 'ssa':
        model = SSA_Net()

    if args.model.arch == 'ms':
        model = MS_Net()

    if args.model.arch == 'msa_':
        model = MSA_Net()
    all_eval_dir = '/home/data1/xiesong/git_repo/removal_thin_cloud/experiments/'+ args.model.arch + '/every_model/'

    list_checkpoint = os.listdir(args.checkpoint_dir)
    for checkpoint in list_checkpoint:
        if int(checkpoint.split('-')[0]) > 3018:
            checkpoint_path = os.path.join(args.checkpoint_dir + checkpoint)
            args.evaluation_dir = os.path.join(all_eval_dir + checkpoint.split('-')[0] + '/')
            os.makedirs(args.evaluation_dir,exist_ok=True)
            print(args.evaluation_dir)
            c_model = create_generator(args, model)

            c_model.load_state_dict(torch.load(checkpoint_path))
            test_model = Testner(c_model)
            test_model .test(args,test_loader)
            test_model.eval(args)

if __name__ == '__main__':

    eval()


from __future__ import absolute_import, division, print_function

import pytest
pytest.importorskip('dask')

from odo.backends.dask import append, Array, merge
from dask.array.core import insert_to_ooc
from dask import core
from odo import convert, into
from odo.utils import tmpfile
import numpy as np
import bcolz

def eq(a, b):
    c = a == b
    if isinstance(c, np.ndarray):
        c = c.all()
    return c

def test_convert():
    x = np.arange(600).reshape((20, 30))
    d = convert(Array, x, blockshape=(4, 5))

    assert isinstance(d, Array)

def test_convert_to_numpy_array():
    x = np.arange(600).reshape((20, 30))
    d = convert(Array, x, blockshape=(4, 5))
    x2 = convert(np.ndarray, d)

    assert eq(x, x2)

def test_append_to_array():
    x = np.arange(600).reshape((20, 30))
    a = into(Array, x, blockshape=(4, 5))
    b = bcolz.zeros(shape=(0, 30), dtype=x.dtype)

    append(b, a)
    assert eq(b[:], x)

    with tmpfile('hdf5') as fn:
        h = into(fn+'::/data', a)
        assert eq(h[:], x)
        h.file.close()

def test_into_inplace():
    x = np.arange(600).reshape((20, 30))
    a = into(Array, x, blockshape=(4, 5))
    b = bcolz.zeros(shape=(20, 30), dtype=x.dtype)

    append(b, a, inplace=True)
    assert eq(b[:], x)

def test_insert_to_ooc():
    x = np.arange(600).reshape((20, 30))
    y = np.empty(shape=x.shape, dtype=x.dtype)
    a = convert(Array, x, blockshape=(4, 5))

    dsk = insert_to_ooc(y, a)
    core.get(merge(dsk, a.dask), list(dsk.keys()))

    assert eq(y, x)

def test__array__():
    x = np.arange(600).reshape((20, 30))
    d = convert(Array, x, blockshape=(4, 5))

    assert eq(x, np.array(d))


import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
from keras.callbacks import Callback

tf.keras.backend.set_floatx('float64')

from multiprocessing import Pool
# from multiprocess import Pool
from keras_pickle_wrapper import KerasPickleWrapper

INPUT_DIM = 4
OUTPUT_DIM = 3

ACTIVATION_DEFAULT = 'sigmoid'
SGD_LR_DEFAULT = 0.01

class WeightsCallback(Callback) :
    # Only works if epochs=1 !
    def __init__(self, epochs, batch_size, train_size, batch_sample_period) :
        self.batch_sample_period = batch_sample_period
        self.n_samples = (epochs * train_size) // (batch_size * batch_sample_period)

        if batch_sample_period == 1 :
            self.on_batch_end = self.on_batch_end_simple
        else :
            self.batch_sample_points = set(range(0, train_size // batch_size, batch_sample_period))
            self.on_batch_end = self.on_batch_end_samples

        self.weights_list = [None] * self.n_samples
        # self.batch_loss = np.empty(self.n_samples)

    def on_batch_end_simple(self, batch, logs=None) :
        self.weights_list[batch] = self.model.get_weights()
        # self.batch_loss[sample_id] = logs['loss']
    def on_batch_end_samples(self, batch, logs=None) :
        if batch in self.batch_sample_points :
            sample_id = batch//self.batch_sample_period
            self.weights_list[sample_id] = self.model.get_weights()
            # self.batch_loss[sample_id] = logs['loss']

l1_norm = lambda m1, m2 : np.sum(np.abs(m1 - m2))

def get_random_init(i, h, o, sigma=1, rng=np.random.default_rng()) :
    dtype = np.float64
    return [
        rng.normal(0, sigma, (i, h)).astype(dtype),
        np.zeros(h).astype(dtype),
        rng.normal(0, sigma, (h, o)).astype(dtype),
        np.zeros(o).astype(dtype)
    ]

def pert_normal(w1, mean=0, std=0.01) :
    w2 = [w.copy() for w in w1]

    w2[0] += np.random.normal(mean, std, w2[0].shape)
    w2[2] += np.random.normal(mean, std, w2[2].shape)

    return w2

def pert_uniform(w1, eps=0.05) :    # Does this actually sample the "ball" uniformly?
    w2 = [w.copy() for w in w1]

    w2[0] += np.random.random_sample(w2[0].shape) * (2 * eps) - eps
    w2[2] += np.random.random_sample(w2[2].shape) * (2 * eps) - eps

    return w2

def flatten_weights(epochs_dict) :
    '''
    Flatten a dict (epochs) of dicts (batches) of lists of arrays (weights) into a list of arrays.
    '''
    batch_weights_flat = []
    for batches in epochs_dict.values() :
        for batch_weights in batches.values() :
            batch_weights_flat.append(batch_weights)

    return batch_weights_flat

def get_dist(w1, w2) :
    # Sum of L1 norm of all matrices in each weights
    return np.sum([l1_norm(u, v) for u, v in zip(w1, w2)])

def get_dists_fromdicts(w1, w2s) :
    '''
    Input is dicts (epochs) of dicts (batches) of weights (i.e. lists of weight matrices).
    '''
    dists = []
    for w2 in w2s :
        dists.append([get_dist(ws1, ws2) for ws1, ws2 in zip(flatten_weights(w1), flatten_weights(w2))])

    return np.array(dists).T    # Convert to plot-friendly format before returning

def get_dists_fromlists(w1, w2s) :
    '''
    Input is list of weights (i.e. lists of weight matrices).
    '''
    dists = []
    for w2 in w2s :
        dists.append([get_dist(ws1, ws2) for ws1, ws2 in zip(w1, w2)])

    return np.array(dists).T    # Convert to plot-friendly format before returning

def get_dists_frommat(w1, w2s) :
    '''
    Input is lists of composite matrices.
    '''
    dists = []
    for w2 in w2s :
        dists.append([l1_norm(u, v) for u, v in zip(w1, w2)])

    return np.array(dists)

def get_dists_dataset_frommat(ws) :
    dists = np.empty((*ws.shape[:2], ws.shape[2]-1, ws.shape[3]))    # *(n_eps, n_init), n_pert, n_sample
    for i_eps in range(dists.shape[0]) :
        for i_init in range(dists.shape[1]) :
            dists[i_eps, i_init] = get_dists_frommat(ws[i_eps, i_init, 0], ws[i_eps, i_init, 1:])

    return dists

def get_permuted_sequences(x_train, y_train, epochs_eff, rng) :
    x_train_eff = np.repeat(x_train, epochs_eff, axis=0)
    y_train_eff = np.repeat(y_train, epochs_eff, axis=0)
    perm = rng.permutation(x_train_eff.shape[0])

    return x_train_eff[perm], y_train_eff[perm]

def get_nonpermuted_sequences(x_train, y_train, epochs_eff) :
    '''
    Do not permute sequences
    '''
    x_train_eff = np.concatenate([x_train] * epochs_eff)
    y_train_eff = np.concatenate([y_train] * epochs_eff)
    # perm = rng.permutation(x_train_eff.shape[0])

    return x_train_eff, y_train_eff

def train(
        xs, ys,
        initial_weights,
        h_dim=16,
        batch_size=1, epochs=1, batch_sample_period=1, n_sample_points=None,
        val_xy=None,
        activation=ACTIVATION_DEFAULT,
        verbose=0,
        shuffle=False,
        return_loss=True,
        return_model=False,
        lr=SGD_LR_DEFAULT,
        with_softmax=False,
        with_sigmoid=False
        ) :

    if n_sample_points is not None :    # Sample points override sample period
        # print(epochs, xs.shape[0], batch_size, n_sample_points)
        batch_sample_period = epochs * xs.shape[0] // batch_size // n_sample_points

    # Infer h_dim from weights
    h_dim = initial_weights[1].size     # Length of first (hidden layer) bias vector

    w_const = keras.initializers.Constant

    model = keras.models.Sequential([
        keras.Input((INPUT_DIM,)),
        keras.layers.Dense(
            h_dim,
            kernel_initializer=w_const(initial_weights[0]),
            bias_initializer=w_const(initial_weights[1]),
            activation=activation
        ),
        keras.layers.Dense(
            OUTPUT_DIM,
            kernel_initializer=w_const(initial_weights[2]),
            bias_initializer=w_const(initial_weights[3]),
            activation=activation
        )
    ])

    if with_softmax :
        model.layers[-1].activation = tf.keras.activations.softmax
    elif with_sigmoid :
        model.layers[-1].activation = tf.keras.activations.sigmoid

    model.compile(
        loss=keras.losses.SparseCategoricalCrossentropy(),
        optimizer=keras.optimizers.legacy.SGD(learning_rate=lr),
        # optimizer='rmsprop',
        # optimizer='adam',
        metrics=[keras.metrics.SparseCategoricalAccuracy()]
    )

    # print(epochs, batch_size, xs.shape[0], batch_sample_period)
    weights_callback = WeightsCallback(epochs, batch_size, xs.shape[0], batch_sample_period)
    # batch_loss = []
    # batch_sample_points = set(range(0, xs.shape[0] // batch_size, batch_sample_period))
    # batch_loss_callback = keras.callbacks.LambdaCallback(
    #     on_batch_end=lambda batch, logs : batch_loss.append(logs['loss'])
    # )

    # model = KerasPickleWrapper(model)

    # history = model().fit(  # Calling the wrapped model to get original
    history = model.fit(
        xs, ys,
        validation_data=val_xy,
        batch_size=batch_size,
        epochs=epochs,
        callbacks=[
            weights_callback,
            # batch_loss_callback,
        ],
        shuffle=shuffle,
        verbose=verbose,
        use_multiprocessing=True
    )

    # Record inital weights and shift batch weights to correspond to loss
    ws = weights_callback.weights_list
    ws.insert(0, initial_weights)
    del ws[-1]

    out_list = [ws]

    if return_loss :    # Compute and return training loss at each snapshot
        loss = get_loss_from_weights_list(ws, model, xs, ys)
        out_list.append(loss)

    if return_model : out_list.append(model)

    if len(out_list) > 1 :      # If mutliple return as tuple
        return tuple(out_list)
    else : return out_list[0]   # If only weights return as object

def train_multiple_seq(n_pert, w1_init, eps, **train_args) :
    '''
    Train an initial condition defined by `w1_init` and `n_pert` perturbations, and return the weights for all.
    '''
    # if 'return_loss' not in train_args :
    #     train_args['return_loss'] = True

    # Train 'reference' instance
    # w1, l1 = train(initial_weights=w1_init, **train_args)
    w1 = train(initial_weights=w1_init, **train_args)

    # Train perturbations
    w2s = []
    # l2s = []
    for i in range(n_pert) :
        # w2, l2 = train(initial_weights=pert_uniform(w1_init, eps), **train_args)
        w2 = train(initial_weights=pert_uniform(w1_init, eps), **train_args)
        w2s.append(w2)
        # l2s.append(l2)

    ws = [w1, *w2s]
    # ls = [l1, *l2s]

    # return ws, ls
    return ws

# Not working — Need to revise
def train_multiple(n_pert, pert_std, n_processes=None, **train_args) :
    xs = train_args.pop('xs')
    ys = train_args.pop('ys')

    h_dim = train_args.pop('h_dim')
    w1_init = get_random_init(INPUT_DIM, h_dim, OUTPUT_DIM)
    w1, bl1 = train(xs, ys, w1_init, **train_args)

    w2s = []
    # h2s = []
    bl2s = []
    with Pool(n_processes) as pool :
        w2_inits = [pert_normal(w1_init, 0, pert_std) for _ in range(n_pert)]
        args = [([xs, ys, w2_inits[i]], train_args) for i in range(n_pert)]
        # print(args)
        res = pool.starmap(train_wrapper, args)

        for w2, bl2 in res :
            w2s.append(w2)
            bl2s.append(bl2)

    # dists = []
    # for w2 in w2s :
    #     dists.append([np.sum([l1_norm(u, v) for u, v in zip(ws1, ws2)]) for ws1, ws2 in zip(flatten_weights(w1), flatten_weights(w2))])

    dists = get_dists_fromdicts(w1, w2s)

    return dists, [bl1, bl2s]

def train_wrapper(args, kwargs) :
    # print(args)
    # print(kwargs)
    return train(*args, **kwargs)

def weights_to_matrix(ws, i, h, o) :
    # ws is a list of 4 matrices
    ws_shape = (i+1 + h+1), (h + o)
    ws_mat = np.zeros(ws_shape)
    ws_mat[:i,:h] = ws[0]
    ws_mat[i,:h] = ws[1]
    ws_mat[i+1:-1,h:] = ws[2]
    ws_mat[-1,h:] = ws[3]

    return ws_mat

def matrix_to_weights(ws_mat, i, h, o) :
    ws = [[], [], [], []]
    ws[0] = ws_mat[:i,:h]
    ws[1] = ws_mat[i,:h]
    ws[2] = ws_mat[i+1:-1,h:]
    ws[3] = ws_mat[-1,h:]

    return ws

def weights_list_to_matrices(ws, i, h, o) :
    ws_shape = (i+1 + h+1), (h + o)
    n_pert = len(ws)
    n_samples = len(ws[0])
    ws_mat = np.zeros((n_pert, n_samples, *ws_shape))

    for i_pert in range(n_pert) :
        for i_pt in range(n_samples) :
            ws_mat[i_pert, i_pt, ...] = weights_to_matrix(ws[i_pert][i_pt], i, h, o)

    return ws_mat

def matrices_to_weights_list(ws_mat, i, h, o) :
    # ws_shape = (i+1 + h+1), (h + o)
    n_pert = ws_mat.shape[0]
    n_samples = ws_mat.shape[1]
    ws = []

    for i_pert in range(n_pert) :
        ws_pert = []
        for i_pt in range(n_samples) :
            ws_pert.append(matrix_to_weights(ws_mat[i_pert, i_pt, ...], i, h, o))
        ws.append(ws_pert)

    return ws

def get_loss_from_weights_list(ws, model, xs, ys) :
    """
    Return `model.loss()` for `xs` and `ys` from each set of weights in `ws`.

    Not a very time efficient solution, but to improve it I would need to
    implement tf training at a lower level...

    Or I could just define my model with a custom loss_tracker ?
    """
    loss = np.empty(len(ws))
    for i, w in enumerate(ws) :
        model.set_weights(w)
        loss[i] = model.loss(ys, model(xs)).numpy()

    return loss

def get_loss_from_matrices(ws, model, xs, ys) :
    """
    Return `model.loss()` for `xs` and `ys` from each weight matrix in `ws`.
    """
    h_dim = ws.shape[-1] - OUTPUT_DIM
    if ws.ndim > 3 :
        ws = matrices_to_weights_list(ws, INPUT_DIM, h_dim, OUTPUT_DIM)   # This could cause confusion if input is for multiple perturbations
    else :
        ws = matrices_to_weights_list(np.expand_dims(ws, axis=0), INPUT_DIM, h_dim, OUTPUT_DIM)[0]

    loss = get_loss_from_weights_list(ws, model, xs, ys)

    return loss

def get_empty_model(h_dim) :
    model = keras.models.Sequential([
        keras.Input((INPUT_DIM,)),
        keras.layers.Dense(
            h_dim,
            activation=ACTIVATION_DEFAULT
        ),
        keras.layers.Dense(
            OUTPUT_DIM,
            activation=ACTIVATION_DEFAULT
        )
    ])

    model.compile(
        loss=keras.losses.SparseCategoricalCrossentropy(),
        optimizer=keras.optimizers.legacy.SGD(learning_rate=SGD_LR_DEFAULT),
        metrics=[keras.metrics.SparseCategoricalAccuracy()]
    )

    return model


#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import annotations
from typing import Dict, Optional
import os
import numpy as np
import torch
import torch.nn as nn
from transformers import BertConfig, BertForSequenceClassification
from .DeepLearning import ConfigDeepLearning

DEFAULT_RANDOM_STATE = None

class ConfigBertClassification(ConfigDeepLearning):
    bert_classification_params = [
        # name, vtype, is_require, default
        ('nict_bert_dir', str, True, None),
        ('num_class', int, True, None),
        ('unique_categories', [list, str], True, None),
        ('random_state', int, False, DEFAULT_RANDOM_STATE),
    ]

    def __init__(
        self: ConfigBertClassification,
        config: Dict,
        config_model_json: str
    ) -> None:
        # model_name
        config['model_name'] = 'BertClassification'
        # load json
        self._load_one(config, config_model_json)
        # init parent class
        self._init_deeplearning(config)
        # set parameters
        for param in self.deeplearning_params:
            self._init_param(config, *param)
        for param in self.bert_classification_params:
            self._init_param(config, *param)
        # value_assertion
        assert(os.path.exists(self.nict_bert_dir))
        return

    def load(self: ConfigBertClassification) -> None:
        return

    def save(self: ConfigBertClassification) -> None:
        return

class BertClassification(nn.Module):
    def __init__(
        self: BertClassification,
        config: Dict,
        config_model_json: str
    ) -> None:
        super().__init__()
        self.config = ConfigBertClassification(config, config_model_json)
        # set random seed
        if self.config.random_state is not None:
            np.random.seed(self.config.random_state)
            torch.manual_seed(self.config.random_state)
            if torch.cuda.is_available():
                torch.cuda.manual_seed(self.config.random_state)
        # load BertForSequenceClassification of NICT BERT
        bert_config = BertConfig.from_pretrained(
            os.path.join(self.config.nict_bert_dir, 'config.json')
        )
        id2label = {i: x for i, x in enumerate(self.config.unique_categories)}
        label2id = {x: i for i, x in enumerate(self.config.unique_categories)}
        bert_config.update({'num_labels': self.config.num_class})
        bert_config.update({'id2label': id2label, 'label2id': label2id})
        if torch.cuda.is_available():
            map_location = torch.device('cuda')
        else:
            map_location = torch.device('cpu')
        state_dict = torch.load(
            os.path.join(self.config.nict_bert_dir, 'pytorch_model.bin'),
            map_location=map_location
        )
        self.model = BertForSequenceClassification.from_pretrained(
            None,
            config=bert_config,
            state_dict=state_dict,
            from_tf=False
        )
        # disable training (changing its weights) of all bert layers
        for name, param in self.model.bert.named_parameters():
            param.requires_grad_(False)
        return

    def forward(
        self: BertClassification,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        token_type_ids: torch.Tensor,
        labels: torch.Tensor
    ) -> torch.Tensor:
        return self.model.forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            labels=labels
        )

    def load(self: BertClassification) -> None:
        bert_config = BertConfig.from_pretrained(
            self.config.config_json
        )
        if torch.cuda.is_available():
            map_location = torch.device('cuda')
        else:
            map_location = torch.device('cpu')
        state_dict = torch.load(
            self.config.binary_path, map_location=map_location
        )
        self.model = BertForSequenceClassification.from_pretrained(
            None,
            config=bert_config,
            state_dict=state_dict,
            from_tf=False
        )
        # disable training (changing its weights) of all bert layers
        for name, param in self.model.bert.named_parameters():
            param.requires_grad_(False)
        return

    def save(self: BertClassification) -> None:
        self.model.save_pretrained(self.config.base_dir)
        return

    def train(self: BertClassification, mode: Optional[bool] = True) -> None:
        self.model.train(mode)
        return

    def eval(self: BertClassification) -> None:
        self.model.eval()
        return


import psycopg2
from psycopg2.extensions import adapt
import logging
import os
import time
import math

logging.basicConfig(level=logging.INFO)

def setup_supabase(vector_name:str, verbose:bool=False):
    hello = f"Setting up supabase database: {vector_name}"
    logging.debug(hello)
    if verbose:
        print(hello)
    setup_database(vector_name, verbose)

def setup_cloudsql(vector_name:str, verbose:bool=False):
    hello = f"Setting up cloudsql database: {vector_name}"
    logging.debug(hello)
    if verbose:
        print(hello)
    setup_database(vector_name, verbose)

def lookup_connection_env(vector_name:str):
    from utils.config import load_config_key

    vs_str = load_config_key("vectorstore", vector_name)
    if vs_str == "supabase":
        return "DB_CONNECTION_STRING"
    elif vs_str == "cloudsql":
        return "PGVECTOR_CONNECTION_STRING"

    raise ValueError("Could not find vectorstore for {vs_str}")

def get_vector_size(vector_name: str):
    from qna.llm import load_config_key
    llm_str = load_config_key("llm", vector_name)

    vector_size = 768
    if llm_str == 'openai':
        vector_size = 1536 # openai

    logging.debug(f'vector size: {vector_size}')
    return vector_size

def setup_database(vector_name:str, verbose:bool=False):

    connection_env = lookup_connection_env(vector_name)

    params = {'vector_name': vector_name, 'vector_size': get_vector_size(vector_name)}

    execute_sql_from_file("sql/sb/setup.sql", params, verbose=verbose, connection_env=connection_env)
    execute_sql_from_file("sql/sb/create_table.sql", params, verbose=verbose, connection_env=connection_env)
    execute_sql_from_file("sql/sb/create_function.sql", params, verbose=verbose, connection_env=connection_env)

    if verbose: print("Ran all setup SQL statements")

    return True

def return_sources_last24(vector_name:str):
    params = {'vector_name': vector_name, 'time_period':'1 day'}
    return execute_sql_from_file("sql/sb/return_sources.sql", params, return_rows=True,
                                 connection_env=lookup_connection_env(vector_name))

def delete_row_from_source(source: str, vector_name:str):
    # adapt the user input and decode from bytes to string to protect against sql injection
    source = adapt(source).getquoted().decode()
    sql_params = {'source_delete': source}
    sql = f"""
        DELETE FROM {vector_name}
        WHERE metadata->>'source' = %(source_delete)s
    """

    do_sql(sql, sql_params=sql_params, connection_env=lookup_connection_env(vector_name))

def do_sql(sql, sql_params=None, return_rows=False, verbose=False, connection_env='DB_CONNECTION_STRING', max_retries=5):

    if connection_env is None:
        raise ValueError("Need to specify connection_env to connect to DB")

    rows = []
    connection_string = os.getenv(connection_env, None)
    if connection_string is None:
        raise ValueError("No connection string")

    for attempt in range(max_retries):
        try:
            connection = psycopg2.connect(connection_string)
            cursor = connection.cursor()

            if verbose:
                logging.info(f"SQL: {sql}")
            else:
                pass
            # execute the SQL - raise the error if already found
            cursor.execute(sql, sql_params)

            # commit the transaction to save changes to the database
            connection.commit()

            if return_rows:
                rows = cursor.fetchall()
            logging.debug("SQL successfully fetched")
            break  # If all operations were successful, break the loop

        except (psycopg2.errors.DuplicateObject,
                psycopg2.errors.DuplicateTable,
                psycopg2.errors.DuplicateFunction) as e:
            logging.debug(str(e))
            if verbose:
                print(str(e))
            continue

        except psycopg2.errors.InternalError as error:
            logging.error(f"InternalError, retrying... Attempt {attempt+1} out of {max_retries}")
            time.sleep(math.pow(2, attempt))  # Exponential backoff
            continue  # Go to the next iteration of the loop to retry the operation

        except (Exception, psycopg2.Error) as error:
            logging.error(f"Error while connecting to PostgreSQL: {str(error)}", exc_info=True)
            continue

        finally:
            if connection:
                cursor.close()
                connection.close()
                logging.debug("PostgreSQL connection is closed")

        # If we've exhausted all retries and still haven't succeeded, raise an error
        if attempt + 1 == max_retries:
            raise Exception("Maximum number of retries exceeded")

    if rows:
        return rows

    return None

def execute_sql_from_file(filepath, params, return_rows=False, verbose=False, connection_env=None):

     # Get the directory of this Python script
    dir_path = os.path.dirname(os.path.realpath(__file__))
    # Build the full filepath by joining the directory with the filename
    filepath = os.path.join(dir_path, filepath)

    # read the SQL file
    with open(filepath, 'r') as file:
        sql = file.read()

    # substitute placeholders in the SQL
    sql = sql.format(**params)
    rows = do_sql(sql, return_rows=return_rows, verbose=verbose, connection_env=connection_env)

    if return_rows:
        if rows is None: return None
        return rows

    return True

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Setup a database",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("vectorname", help="The namespace for vectorstore")
    parser.add_argument("connection_env", help="The connection environment string", default="DB_CONNECTION_STRING")

    args = parser.parse_args()
    config = vars(args)

    vector_name = config.get('vectorname', None)
    if vector_name is None:
        raise ValueError("Must provide a vectorname")

    setup_database(vector_name, verbose=True, connection_env=config.get("connection_env"))


from typing import List
from fastapi import APIRouter, Depends, HTTPException, Request
from common.exceptions import NotFoundException, InconsistentPipelineException, ConflictException
from pipelines.service import PipelinesService
from common.query_parameters import QueryParameters
from pipelines.models import PipelineRead, PipelineUpdate, PipelineCreate, Pipeline, \
    PipelineReadWithPipelineStepsAndTasks, PipelinesWithCount
from pipeline_steps.models import PipelineStep
from uuid import UUID, uuid4
from sqlalchemy.exc import CompileError

router = APIRouter()

@router.get(
    "/pipelines/{pipeline_id}",
    summary="Get one pipeline",
    responses={
        404: {"detail": "Pipeline Not Found"},
        400: {"detail": "Bad Request"},
        500: {"detail": "Internal Server Error"},
    },
    response_model=PipelineReadWithPipelineStepsAndTasks,
)
def get_one(
        pipeline_id: UUID,
        pipelines_service: PipelinesService = Depends()
):
    pipeline = pipelines_service.find_one(pipeline_id)
    if not pipeline:
        raise HTTPException(status_code=404, detail="Pipeline Not Found")

    return pipeline

@router.get(
    "/pipelines",
    summary="Get many pipelines",
    response_model=PipelinesWithCount | List[PipelineRead],
)
def get_many_pipelines(
        with_count: bool = False,
        query_parameters: QueryParameters = Depends(),
        pipelines_service: PipelinesService = Depends(),
):
    try:
        if query_parameters.search:
            query_parameters.search = query_parameters.search.lower()

        if with_count:
            count, pipelines = pipelines_service.find_many_with_total_count(
                query_parameters.search,
                query_parameters.skip,
                query_parameters.limit,
                query_parameters.order_by,
                query_parameters.order,
                query_parameters.tags,
                query_parameters.status
            )

            return PipelinesWithCount(count=count, pipelines=pipelines)
        else:
            pipelines = pipelines_service.find_many(
                query_parameters.search,
                query_parameters.skip,
                query_parameters.limit,
                query_parameters.order_by,
                query_parameters.order,
                query_parameters.tags,
                query_parameters.status
            )

            return pipelines

    except CompileError as e:
        raise HTTPException(status_code=400, detail=str(e))

@router.post(
    "/pipelines",
    summary="Create a pipeline",
    responses={
        400: {"detail": "Bad Request"},
        500: {"detail": "Internal Server Error"},
    },
    response_model=PipelineReadWithPipelineStepsAndTasks,
)
def create(
        request: Request,
        pipeline: PipelineCreate,
        pipelines_service: PipelinesService = Depends(),
):
    try:
        pipeline = pipelines_service.create(pipeline, request.app)
    except InconsistentPipelineException as e:
        raise HTTPException(status_code=400, detail=str(e))
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=str(e))
    except ConflictException as e:
        raise HTTPException(status_code=409, detail=str(e))

    return pipeline

@router.patch(
    "/pipelines/{pipeline_id}",
    summary="Update a pipeline",
    responses={
        404: {"detail": "Pipeline Not Found"},
        500: {"detail": "Internal Server Error"},
    },
    response_model=PipelineReadWithPipelineStepsAndTasks,
)
def update(
        request: Request,
        pipeline_id: UUID,
        pipeline_update: PipelineUpdate,
        pipelines_service: PipelinesService = Depends(),
):
    try:
        pipeline = pipelines_service.update(request.app, pipeline_id, pipeline_update)
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=str(e))
    except InconsistentPipelineException as e:
        raise HTTPException(status_code=400, detail=str(e))
    except ConflictException as e:
        raise HTTPException(status_code=409, detail=str(e))

    return pipeline

@router.delete(
    "/pipelines/{pipeline_id}",
    summary="Delete a pipeline",
    responses={
        204: {"detail": "Pipeline Removed"},
        404: {"detail": "Pipeline Not Found"},
        500: {"detail": "Internal Server Error"},
    },
    status_code=204
)
def delete(
        request: Request,
        pipeline_id: UUID,
        pipelines_service: PipelinesService = Depends(),
):
    try:
        pipelines_service.delete(pipeline_id, request.app)
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=str(e))

@router.post(
    "/pipeline/check",
    summary="Check if a pipeline is valid",
    responses={
        200: {"valid": True},
        400: {"detail": "Bad Request"},
        500: {"detail": "Internal Server Error"},
    }
)
def check(
        pipeline: PipelineCreate,
        pipelines_service: PipelinesService = Depends(),
):
    try:
        new_pipeline_id = uuid4()
        new_steps = []
        for step in pipeline.steps:
            new_step = PipelineStep(
                identifier=step.identifier,
                needs=step.needs,
                condition=step.condition,
                inputs=step.inputs,
                service_id=step.service_id,
                pipeline_id=new_pipeline_id,
            )
            new_steps.append(new_step)

        new_pipeline = Pipeline(
            id=new_pipeline_id,
            name=pipeline.name,
            description=pipeline.description,
            summary=pipeline.summary,
            steps=new_steps,
            data_in_fields=pipeline.data_in_fields,
            data_out_fields=pipeline.data_out_fields,
        )

        if pipelines_service.check_pipeline_consistency(new_pipeline):
            return {"valid": True}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


# -*- coding: utf-8 -*-
__author__ = 'palmtale'

from pyramid.i18n import TranslationStringFactory

_ = TranslationStringFactory("index")

title = _("title", "Palms Main")
com_name_first = _("com_name_first", "Palms")
com_name_last = _("com_name_last", "Com")
sunday = _("sunday", "Sunday")
monday = _("monday", "Monday")
tuesday = _("tuesday", "Tuesday")
wednesday = _("wednesday", "Wednesday")
thursday = _("thursday", "Thursday")
friday = _("friday", "Friday")
saturday = _("saturday", "Saturday")
duty_title = _("duty_title", "Duty time situation")
duty_time_label = _("duty_time_label", "Duty Time")
duty_yesterday_label = _("duty_yesterday_label", "Yesterday")
duty_check_in = _("duty_check_in", "Check in")
duty_check_all = _("duty_check_all", "Check all")
content_title = _("content_title", "Home")
content_subtitle = _("content_subtitle", "Welcome to Palm Com management panel")
site_map_prompt = _("site_map_prompt", "You are here:")
site_map_index = _("site_map_index", "PalmCom")
site_map_home = _("site_map_home", "Home")


import numpy as np
import matplotlib.pyplot as plt
from math import *
# import ROOT
import matplotlib.patches as mpatches
import statistics
from copy import deepcopy
import sys
import os
import random
import yaml
import matplotlib
import glob
from grepfunc import grep_iter
import re
import json
import argparse
from termcolor import colored
import mplhep as hep
from matplotlib.backends.backend_pdf import PdfPages
from mpl_toolkits.axes_grid1 import AxesGrid
from scipy.stats import sem

regex_typelabel=re.compile("Q")
regex_amodule=re.compile("dPosXYZ")
regex_rmodule=re.compile("dRotXYZ")
labels=["Tx","Ty","Tz","Rx","Ry","Rz"]
positions=["x_global","y_global","z_global"]
trackInfo=["nTracks","nHits"]
stations = ["T1", "T2", "T3"]
layers = ["U", "V", "X1", "X2"]

colors = ['black', 'blue', 'red', 'green', 'yellow', 'magenta', 'brown', 'cyan']
markers = ['o', 'x', 'd', 'D', '.', 'v', 's', 'p']

def open_alignment(thisfile,convergence=True):
    with open(thisfile) as f:
        align_output=json.load(f)

    convergences=align_output.pop("converged")

    for alignable in align_output.keys():
        for label in labels+positions+trackInfo:
            if "FT" in alignable:
                align_output[alignable][label]=[float(ele.strip(',')) for ele in align_output[alignable][label]]

    if convergence:
        align_output["convergence"]=convergences
    return align_output

def convertGlobal(align_external,halfmoduleAlignables):
    align_output=deepcopy(align_external)
    for alignable in halfmoduleAlignables:
        # rule for readout flip in x per layer:
        if "Quarter0" in alignable or "Quarter1" in alignable or "Q0" in alignable or "Q1" in alignable:
            for label in ["Tx","Rx","Ry","Ty"]:
                if label in align_output[alignable].keys():
                    align_output[alignable][label]=[-value for value in align_output[alignable][label]]
            # minus sign on x and Rx and y and Ry for all iterations
        if "X1" in alignable or "V" in alignable:
            # minus sign on x and Rx and z and Rz for all iterations
            for label in ["Tx","Rx","Rz","Tz"]:
                if label in align_output[alignable].keys():
                    align_output[alignable][label]=[-value for value in align_output[alignable][label]]

    return align_output

def get_unc(input):
    n_files = len(input)
    Tx_chi2 = [[] for _ in range(n_files)]
    Ty_chi2 = [[] for _ in range(n_files)]
    Tz_chi2 = [[] for _ in range(n_files)]

    Rx_chi2 = [[] for _ in range(n_files)]
    Ry_chi2 = [[] for _ in range(n_files)]
    Rz_chi2 = [[] for _ in range(n_files)]
    for i in range(n_files):
        Tx_chi2[i].append(input[i]['Tx_chi2'][0])
        Ty_chi2[i].append(input[i]['Ty_chi2'][0])
        Tz_chi2[i].append(input[i]['Tz_chi2'][0])

        Rx_chi2[i].append(input[i]['Rx_chi2'][0])
        Ry_chi2[i].append(input[i]['Ry_chi2'][0])
        Rz_chi2[i].append(input[i]['Rz_chi2'][0])

    return np.array(Tx_chi2), np.array(Ty_chi2), np.array(Tz_chi2), np.array(Rx_chi2), np.array(Ry_chi2), np.array(Rz_chi2)

files_Tx = [\
         "txt_files/parsed_out/parsedlog_269045_Tx_0_2mm.json",
         "txt_files/parsed_out/parsedlog_269045_Tx_0_4mm.json",
         "txt_files/parsed_out/parsedlog_269045_Tx_0_6mm.json",
         "txt_files/parsed_out/parsedlog_269045_Tx_0_8mm.json",
         "txt_files/parsed_out/parsedlog_269045_base_noSurvey.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_2mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_3mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_4mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_5mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_6mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_7mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_8mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_9mm_noS.json",
         "txt_files/parsed_out/Tx/parsedlog_269045_Tx_10mm_noS.json",
]
files_Ty = [\
         "txt_files/parsed_out/parsedlog_269045_base_noSurvey.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_2mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_3mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_4mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_5mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_6mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_7mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_8mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_9mm_noS.json",
         "txt_files/parsed_out/Ty/parsedlog_269045_Ty_10mm_noS.json",
]
files_Rx = [\
        "txt_files/parsed_out/parsedlog_269045_base_noSurvey.json",
        "txt_files/parsed_out/Rx/parsedlog_Rx_03mrad.json",
        "txt_files/parsed_out/Rx/parsedlog_Rx_04mrad.json",
        "txt_files/parsed_out/Rx/parsedlog_Rx_05mrad.json"
]
files_Ry = [\
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_4mrad.json", # 0.0004
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_3mrad.json", # 0.0003
        "txt_files/parsed_out/parsedlog_269045_base_noSurvey.json", # 0.0002 base
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_1mrad.json", # 0.0001
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_02mrad.json", # 0.00002
        "txt_files/parsed_out/Ry/parsedlog_Ry_3microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_2microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_1microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_8microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_6microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_5microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_4microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_44microrad.json",
        "txt_files/parsed_out/Ry/parsedlog_Ry_0_2microrad.json",
]

legendlabels_Tx=[\
        "0_2mm",
        "0_4mm",
        "0_6mm",
        "0_8mm",
        "base",
        "2mm",
        "3mm",
        "4mm",
        "5mm",
        "6mm",
        "7mm",
        "8mm",
        "9mm",
        "0mm"
]
legendlabels_Ty = [\
        "base",
        "2mm",
        "3mm",
        "4mm",
        "5mm",
        "6mm",
        "7mm",
        "8mm",
        "9mm",
        "10mm",
]
legendlabels_Rx = [\
        "base",
        "03mrad",
        "04mrad",
        "05mrad",
]

legendlabels_Ry = [\
        "04mrad",
        "03mrad",
        "base",
        "01mrad",
        "0_02mrad",
        "3microRad",
        "2microRad",
        "1microRad",
        "08microRad",
        "06microRad",
        "05microRad",
        "0_44microRad",
        "04microRad",
        "02microRad",
]

def get_chi2_values(input_files):
    outputs=[open_alignment(thisfile) for thisfile in input_files]
    plotted_alignables=[]
    for align_block in outputs:
        thislist=[]
        for key in align_block.keys():
            thislist.append(key)
        plotted_alignables.append(thislist)
    outputs=[convertGlobal(align_block,plotted_alignables[0]) for align_block in outputs]

    Tx_chi2_tx = get_unc(outputs)[0]
    Ty_chi2_tx = get_unc(outputs)[1]
    Tz_chi2_tx = get_unc(outputs)[2]
    Rx_chi2_tx = get_unc(outputs)[3]
    Ry_chi2_tx = get_unc(outputs)[4]
    Rz_chi2_tx = get_unc(outputs)[5]

    return Tx_chi2_tx, Ty_chi2_tx, Tz_chi2_tx, Rx_chi2_tx, Ry_chi2_tx, Rz_chi2_tx

chi2_values_from_Tx_changes = get_chi2_values(files_Tx)
chi2_values_from_Ty_changes = get_chi2_values(files_Ty)
chi2_values_from_Rx_changes = get_chi2_values(files_Rx)
chi2_values_from_Ry_changes = get_chi2_values(files_Ry)

Tz_chi2_vals = [2586.76]

# 7mm and 8mm are missing, last entries are 5mm, 6mm, 7mm, 10mm
Tx_chi2_tz = [15692/768, 8182.69/768, 5278.57/768, 5770.66/768, 12196.7/768, 9622.9/768, 10590.6/768]
Ty_chi2_tz = [1127.64/768, 1122.31/768, 1119.17/768, 1120.25/768, 1125.93/768, 1125.58/768, 1126.06/768]
Tz_chi2_tz = [646.69/768, 287.418/768, 161.672/768, 103.47/768, 71.8544/768, 52.791/768, 25.8676/768]
Rx_chi2_tz = [3086.38/768, 3086.38/768, 3086.38/768, 3086.38/768, 3086.38/768, 3086.38/768, 3086.38/768]
Ry_chi2_tz = [0.00366032/768, 0.0036605/768, 0.00369746/768, 0.00370069/768, 0.0036554/768, 0.00368278/768, 0.0036503/768]
Rz_chi2_tz = [708.737/768, 679.203/768, 658.385/768, 658.059/768, 697.696/768, 678.97/768,  687.408/768]

Tz_data = np.array([Tx_chi2_tz, Ty_chi2_tz, Tz_chi2_tz, Rx_chi2_tz, Ry_chi2_tz, Rz_chi2_tz])
print(Tz_data)
correct_order = [2, 0, 1, 3, 6, 4, 5, 7, 10, 8, 9, 11]
Tx_err = [0.2, 0.4, 0.6, 0.8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]  # microns
Ty_err = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] # microns
Tz_err = [2,3,4,5,6,7,10] # microns

Rx_err = [0.2, 0.3, 0.4, 0.5] # microrad
Ry_err = [0.4, 0.3, 0.2, 0.1, 0.02, 0.003, 0.002, 0.001, 0.0008, 0.0006, 0.0005, 0.00044, 0.0004, 0.0002] # microrad

total_n_dofs = 768

def plotting(x_range, data, dofs, label):
    if label == 'Ry':
        x_vals = np.linspace(0, 14, 14)
        print(f"label: {label}")
        print(data[0])
        plt.plot(x_vals, data[0] / dofs, marker='.', linestyle='--', label='Tx_chi')
        plt.plot(x_vals, data[1] / dofs, marker='.', linestyle='--', label='Ty_chi')
        plt.plot(x_vals, data[2] / dofs, marker='.', linestyle='--', label='Tz_chi')
        plt.plot(x_vals, data[3] / dofs, marker='.', linestyle='--', label='Rx_chi')
        plt.plot(x_vals, data[4] / dofs, marker='.', linestyle='--', label='Ry_chi')
        plt.plot(x_vals, data[5] / dofs, marker='.', linestyle='--', label='Rz_chi')
        plt.xticks(x_vals, x_range, rotation=45)
        plt.xlabel('[mu rad]')
    if label == 'Tz':
        print(f"label: {label}")
        print(data[0])
        plt.plot(x_range, data[0], marker='.', linestyle='--', label='Tx_chi')
        plt.plot(x_range, data[1], marker='.', linestyle='--', label='Ty_chi')
        plt.plot(x_range, data[2], marker='.', linestyle='--', label='Tz_chi')
        plt.plot(x_range, data[3], marker='.', linestyle='--', label='Rx_chi')
        plt.plot(x_range, data[4], marker='.', linestyle='--', label='Ry_chi')
        plt.plot(x_range, data[5], marker='.', linestyle='--', label='Rz_chi')
    if label in ['Tx', 'Ty']:
        print(f'else: {label}')
        print(data)
        plt.plot(x_range, data[0] / dofs, marker='.', linestyle='--', label='Tx_chi')
        plt.plot(x_range, data[1] / dofs, marker='.', linestyle='--', label='Ty_chi')
        plt.plot(x_range, data[2] / dofs, marker='.', linestyle='--', label='Tz_chi')
        plt.plot(x_range, data[3] / dofs, marker='.', linestyle='--', label='Rx_chi')
        plt.plot(x_range, data[4] / dofs, marker='.', linestyle='--', label='Ry_chi')
        plt.plot(x_range, data[5] / dofs, marker='.', linestyle='--', label='Rz_chi')
        plt.xlabel('[mu m]')
    if label == 'Rx':
        print(f'else: {label}')
        print(data)
        plt.plot(x_range, data[0] / dofs, marker='.', linestyle='--', label='Tx_chi')
        plt.plot(x_range, data[1] / dofs, marker='.', linestyle='--', label='Ty_chi')
        plt.plot(x_range, data[2] / dofs, marker='.', linestyle='--', label='Tz_chi')
        plt.plot(x_range, data[3] / dofs, marker='.', linestyle='--', label='Rx_chi')
        plt.plot(x_range, data[4] / dofs, marker='.', linestyle='--', label='Ry_chi')
        plt.plot(x_range, data[5] / dofs, marker='.', linestyle='--', label='Rz_chi')
        plt.xlabel('[mu rad]')
    plt.hlines(1, min(x_range), max(x_range), 'black')
    plt.legend()
    plt.grid()
    plt.title(f'chi2 / dof changing only {label} error')
    plt.ylabel('chi2 per dof')
    # plt.show()
    plt.savefig(f'chi2_plots/{label}_out.pdf')
    plt.clf()

plotting(Tx_err, chi2_values_from_Tx_changes, total_n_dofs, 'Tx')
plotting(Ty_err, chi2_values_from_Ty_changes, total_n_dofs, 'Ty')
plotting(Tz_err, Tz_data, total_n_dofs, 'Tz')
plotting(Rx_err, chi2_values_from_Rx_changes, total_n_dofs, 'Rx')
plotting(Ry_err, chi2_values_from_Ry_changes, total_n_dofs, 'Ry')

errors_v1 = [0.001, 0.0015, 0.004, 0.0002, 0.0002, 0.0002]
errors_v2 = [0.0005, 0.002, 0.002, 0.0002, 0.0002, 0.0002]
errors_v3 = [0.0002, 0.002, 0.002, 0.0002, 0.0002, 0.0002]
errors_v4 = [0.00018, 0.0015, 0.0018, 0.0004, 0.00000044, 0.00019]

from scipy.stats import linregress
x1 = Tx_err[0:7]
y1 = (chi2_values_from_Tx_changes[0][0:7] / total_n_dofs).T[0]
print(x1, y1)
erg = linregress(x1, y1)
print(erg)

# try curve fit
from scipy.optimize import curve_fit
def func(x, a, b, c):
    return a * np.exp(-b * x) + c
def Gauss(x, A, B):
    y = A*np.exp(-1*B*x**2)
    return y
x2 = Tx_err[6:]
y2 = (chi2_values_from_Tx_changes[0][6:] / total_n_dofs).T[0]

xfull = np.array(Tx_err)
yfull = np.array((chi2_values_from_Tx_changes[0] / total_n_dofs).T[0])

xdata1 = np.array(x1)
ydata1 = np.array(y1)

xdata2 = np.array(x2)
ydata2 = np.array(y2)

plt.plot(xfull, yfull, 'k-', label='data')
popt1, pcov1 = curve_fit(func, xdata1, ydata1)
popt2, pcov2 = curve_fit(func, xdata2, ydata2)
plt.plot(xdata1, func(xdata1, *popt1), 'r-',label='fit: a=%5.3f, b=%5.3f, c=%5.3f' % tuple(popt1))
# plt.plot(xdata2, func(xdata2, *popt2), 'b-',label='fit: A=%5.3f, B=%5.3f, c=%5.3f' % tuple(popt2))
plt.axvline(x = 0.29, color = 'g', linestyle='dashed', label='Tx unc. = 0.29 micron')
plt.axhline(y = 1, color = 'y', linestyle='dashed', label='chi2 / dof = 1')
plt.grid()
plt.legend()
plt.xlabel('Tx unc [micron]')
plt.ylabel('chi2 / dof')
plt.title('Tx uncertainty tuning')
plt.savefig('fit_Tx.pdf')
plt.show()
# print(func(xdata, *popt))

# FIXME:
# 1. correct the x label to be LaTeX formated
# 2. Tx: more alignments between 0 and 1 microns, also fit a function (linear, polynomial 2nd order and find chi2 / dof = 1)
# 3. Ty: more alignments between 1.2 and 1.7 microns
# 4. Tz: more alignments between 1 and 2 microns
# 5. Rx: more alignments around 0.4 micro rad (0.0004 milli rad)
# 6. Ry: more alignments roughly at 0.00044 microns
# 7: Rz: please also do a plot for Rz around 0.0002 milli rad


import imports
import speak_take

def TaskExe():
    speak_take.speak("Hello, I am jarvis!")
    speak_take.speak("How can i help you? ")

    def GoogleSearch(term):
        query = term.replace("jarvis","")
        query = query.replace("what is","")
        query = query.replace("how to","")
        query = query.replace("who is","")
        query = query.replace("what do we mean by","")
        query = query.replace("what do you mean by","")
        query = query.replace(" ","")

        Query = str(term)
        imports.pywhatkit.search(Query)
        if "google search" in Query:
            search = imports.wikipedia.summary(Query,2)
            speak_take.speak(f":According to your search : {search}")

        else:
            pass

    def speedtest():
        import speedtest
        speak_take.speak("Checking Speed")
        speed = speedtest.Speedtest()
        downloading = speed.download()
        correctDown= int(downloading/800000)
        uploading= speed.upload()
        correctUpload=int(uploading/800000)

        if 'uploading' in query:
            speak_take.speak(f"The uploading speed is {correctUpload} mbps")

        elif 'downloading' in query:
            speak_take.speak(f"The downloading speed is {correctDown} mbps")

        else:
            speak_take.speak(f"The downloading speed is {correctDown} mbp/s and uploading speed is {correctUpload} mbps")

    def takeHindi():
        command = imports.sr.Recognizer()
        with imports.sr.Microphone() as source:
            print("Listening.............")
            command.pause_threshold = 1
            audio = command.listen(source)

            try:
                print("Recognizing............")
                query = command.recognize_google(audio,language='en-in')
                print(f'you said : {query}')

            except:
                return "none"

            return query.lower()

    def Tran():
        speak_take.speak('Tell me the Line!')
        line = takeHindi()
        traslate = imports.Translator()
        result = traslate.translate(line)
        Text = result.text
        speak_take.speak(f"The Translation Of This Line Is:",Text)

    def Temp():
        search="temperature in pune"
        url = f"https://www.google.com/search?q={search}"
        r=imports.requests.get(url)
        data= imports.BeautifulSoup(r.text,"html.parser")
        temperature=data.find("div",class_="BNeawe").text
        speak_take.speak(f"the temperature outside is {temperature}")

    def music():
        speak_take.speak("Tell me the name of the song!")
        musicname = speak_take.takecommand()
        imports.pywhatkit.playonyt(musicname)
        speak_take.speak("Your song has been started!, Enjoy sir")

    def whatsapp():
        speak_take.speak("Tell me the name of the person")
        name = speak_take.takecommand()

        if 'abc' in name:
            speak_take.speak("Tell me the message")
            msg=speak_take.takecommand()
            speak_take.speak("tell me the time sir")
            speak_take.speak("time in hour")
            hour=int(speak_take.takecommand())
            speak_take.speak("time in minutes")
            min=int(speak_take.takecommand())
            imports.pywhatkit.sendwhatmsg("+919307665461",msg,hour,min,20)
            speak_take.speak("ok sir, sending whatsapp message")

        elif 'about' in name:
            speak_take.speak("Tell me the message")
            msg=speak_take.takecommand()
            speak_take.speak("tell me the time sir")
            speak_take.speak("time in hour")
            hour=int(speak_take.takecommand())
            speak_take.speak("time in minutes")
            min=int(speak_take.takecommand())
            imports.pywhatkit.sendwhatmsg("+919730863521",msg,hour,min,20)
            speak_take.speak("ok sir, sending whatsapp message")

        else:
            speak_take.speak("tell me the phone number")
            phone=int(speak_take.takecommand())
            ph = "+91" + phone
            speak_take.speak("Tell me the message")
            msg=speak_take.takecommand()
            speak_take.speak("tell me the time sir")
            speak_take.speak("time in hour")
            hour=int(speak_take.takecommand())
            speak_take.speak("time in minutes")
            min=int(speak_take.takecommand())
            imports.pywhatkit.sendwhatmsg(ph,msg,hour,min,20)
            speak_take.speak("ok sir, sending whatsapp message")

    def ChromeAuto():
        speak_take.speak('Chrome Automation Started!!')
        command = speak_take.takecommand()

        if 'close this tab' in command:
            imports.keyboard.press_and_release('ctrl + w')

        elif 'open new tab' in command:
            imports.keyboard.press_and_release('ctrl + t')

        elif 'open new window' in command:
            imports.keyboard.press_and_release('ctrl + n')

        elif 'history' in command:
            imports.keyboard.press_and_release('ctrl + h')

    def screenshot():
        speak_take.speak("Ok sir! What should i name that file? ")
        path = speak_take.takecommand()
        path1name  = path + ".png"
        path1 = "C:\\Users\\Ketan\\OneDrive\\Pictures\\Screenshots\\" +path1name
        kk = imports.pyautogui.screenshot()
        kk.save(path1)
        imports.os.startfile("C:\\Users\\Ketan\\OneDrive\\Pictures\\Screenshots")
        speak_take.speak("Here is your screenshot")

    def Dict():
        speak_take.speak('Dictionary Activated !!')
        speak_take.speak('Tell me the Word Sir...!')
        prob1 = speak_take.takecommand()

        if 'meaning' in prob1:
            # probl = probl.replace("what is the","")
            # probl = probl.replace("jarvis","")
            # probl = probl.replace("of","")
            probl = probl.replace("meaning of","")
            result = imports.diction.meaning(probl)
            speak_take.speak(f'The Meaning For {probl} is {result}')

        else:
            speak_take.speak("tejam")

    def openApps():
        speak_take.speak("Ok sir, Wait a second!")

        if "code" in query:
            imports.os.startfile("C:\\Users\\Ketan\\AppData\\Local\\Programs\\Microsoft VS Code\\Code.exe")

        elif "brave" in query:
            imports.os.startfile("C:\\Program Files\\BraveSoftware\\Brave-Browser\\Application\\brave.exe")

        elif "whatsapp" in query:
            imports.os.startfile("C:\\Users\\Ketan\\AppData\\Local\\WhatsApp\\WhatsApp.exe")

        elif "chrome" in query:
            imports.os.startfile("C:\\Program Files\\Google\\Chrome\\Application\\chrome.exe")

        elif "facebook" in query:
            imports.webbrowser.open("https://www.facebook.com/")

        elif "youtube" in query:
            imports.webbrowser.open("https://www.youtube.com/")

        elif "instagram" in query:
            imports.webbrowser.open("https://www.instagram.com/")

        elif "maps" in query:
            imports.webbrowser.open("https://www.google.com/maps/@18.6793851,73.8536679,15z")

    def closeApps():
        speak_take.speak("Ok sir!, Wait a second..")

        if "chrome" in query:
            imports.os.system("TASKKILL /F /im chrome.exe")

        elif "brave" in query:
            imports.os.system("TASKKILL /F /im brave.exe")

        elif "code" in query:
            imports.os.system("TASKKILL /F /im Code.exe")

        elif "whatsapp" in query:
            imports.os.system("TASKKILL /F /im WhatsApp.exe")

    def YoutubeSearch(term):
        result = "https://www.youtube.com/results?search_query=" + term
        imports.webbrowser.open(result)
        speak_take.speak("This is what i found for your search")
        imports.pywhatkit.playonyt(term)
        speak_take.speak("This may also help you Sir")

    def YoutubeAuto():
        speak_take.speak("Whats your Command ?")
        comm = speak_take.takecommand()

        if 'pause' in comm:
            imports.keyboard.press('space bar')

        elif 'restart' in comm:
            imports.keyboard.press('0')

        elif 'mute' in comm:
            imports.keyboard.press('m')

        elif 'skip' in comm:
            imports.keyboard.press('l')

        elif 'back' in comm:
            imports.keyboard.press('j')

        elif 'full screen' in comm:
            imports.keyboard.press('f')

        elif 'film mode' in comm:
            imports.keyboard.press('t')

        speak_take.speak('Done Sir')

    while True:
        query=speak_take.takecommand()

        if "hello" in query:
            speak_take.speak("hello sir i am jarvis")
            speak_take.speak("your personal assistant")
            speak_take.speak("How may i help you?")

        elif "how are you" in query:
            speak_take.speak("I am fine sir")
            speak_take.speak("what about you")

        elif "you need a break" in query:
            speak_take.speak("ok sir")
            speak_take.speak("you can call me anytime")
            speak_take.speak("Just say Wake Up jarvis!")
            break

        elif "bye" in query:
            speak_take.speak("bye sir")
            speak_take.speak("have a good day")
            break

        elif "youtube search" in query:
            query = query.replace("youtube search","")
            query = query.replace("jarvis","")
            YoutubeSearch(query)

        elif "google search" in query:
            speak_take.speak("This is what i found for your search sir!")
            query = query.replace("google search","")
            query = query.replace("jarvis","")
            # pywhatkit.search(query)
            GoogleSearch(query)
            speak_take.speak("Done Sir!!")

        elif "website" in query:
            speak_take.speak("Ok sir, Launching")
            query = query.replace("website","")
            query = query.replace("jarvis","")
            query = query.replace(" ","")
            web1 = query.replace("open","")
            web2 = "https://www." + web1 + ".com"
            imports.webbrowser.open(web2)
            speak_take.speak("launched sir!")

        elif "launch" in query:
            speak_take.speak("Tell me the name of the website!")
            name = speak_take.takecommand()
            web = 'https://www.' + name + '.com'
            imports.webbrowser.open(web)
            speak_take.speak("Done Sir!")

        # elif "whatsapp message" in query:
        #     whatsapp()

        elif "play song" in query:
            music()

        elif "wikipedia" in query:
            speak_take.speak("Searching wikipedia")
            query=query.replace("jarvis","")
            query=query.replace("wikipedia","")
            wiki= imports.wikipedia.summary(query,2)
            speak_take.speak(f"According to Wikipedia :{wiki}")

        elif "screenshot" in query:
            screenshot()

        elif "open facebook" in query:
            openApps()

        elif "open instagram" in query:
            openApps()

        elif "open code" in query:
            openApps()

        elif "open youtube" in query:
            openApps()

        elif "open brave" in query:
            openApps()

        elif "open chrome" in query:
            openApps()

        elif "open map" in query:
            openApps()

        elif "open whatsapp" in query:
            openApps()

        elif "repeat my words" in query:
            speak_take.speak("speak sir")
            ab= speak_take.takecommand()
            speak_take.speak(f"you said : {ab}")

        elif "close chrome" in query:
            closeApps()

        elif "close brave" in query:
            closeApps()

        elif "close whatsapp" in query:
            closeApps()

        elif "close code" in query:
            closeApps()

        elif "chrome automation" in query:
            ChromeAuto()

        elif 'stop' in query:
            imports.keyboard.press('spacebar')

        elif 'restart' in query:
            imports.keyboard.press('0')

        elif 'mute' in query:
            imports.keyboard.press('m')

        elif 'skip' in query:
            imports.keyboard.press('l')

        elif 'back' in query:
            imports.keyboard.press('j')

        elif 'full screen' in query:
            imports.keyboard.press('f')

        elif 'film mode' in query:
            imports.keyboard.press('t')

        elif 'pause' in query:
            imports.keyboard.press('k')

        elif 'youtube tool' in query:
            YoutubeAuto()

        elif 'dictionary' in query:
            Dict()

        elif 'joke' in query:
            get = imports.pyjokes.get_joke()
            speak_take.speak(get)

        elif "that was pretty cringe" in query:
            speak_take.speak("reaaaalllllly??????????????????")

        elif 'alarm' in query:
            speak_take.speak("enter the time")
            time=input(": Enter the time :")

            while True:
                Time_Ac= imports.datetime.datetime.now()
                now=Time_Ac.strftime("%H:%M:%S")

                if now==time:
                    speak_take.speak("time to wake up sir")
                    imports.playsound('D:\\jarvis project\\newsound.mp3')
                    speak_take.speak("Alarm closed")

                elif now>time:
                    break

        elif "remember that" in query:
            rememberMsg = query.replace("remember that","")
            rememberMsg = query .replace("jarvis","")
            speak_take.speak("You Told me to remind you that :" +rememberMsg)
            remember = open("data.txt","w")
            remember.write(rememberMsg)
            remember.close()

        elif "what do you remember" in query:
            remember = open("data.txt","r")
            speak_take.speak("You told me to remind that" +remember.read())

        elif 'please search' in query:
            import wikipedia as googleScrap
            query = query.replace("jarvis","")
            query = query.replace("google search","")
            query = query.replace("google","")
            speak_take.speak("This is what i found on the web")

            try:
                # pywhatkit.search(query)
                result = googleScrap.summary(query,2)
                speak_take.speak(result)

            except:
                speak_take.speak("no speakable data available")

        elif 'temperature' in query:
            Temp()

        # elif 'translator' in query:
        #     Tran()

        elif 'downloading speed' in query:
            speedtest()

        elif 'uploading speed' in query:
            speedtest()

        elif 'internet speed' in query:
            speedtest()

        elif 'whatsapp message' in query:
            query = query.replace("jarvis","")
            query = query.replace("send","")
            query = query.replace("whatsapp message","")
            query = query.replace("to","")
            name = query

            if 'laptop' in name:
                numb = "9067084690"
                speak_take.speak(f"Whats the Message For {name}")
                mess = speak_take.takecommand()
                whatsapp.whatsapp(numb,mess)

            elif 'pretty' in name:
                numb = "9307665461"
                speak_take.speak(f"Whats the Message For {name}")
                mess = speak_take.takecommand()
                whatsapp.whatsapp(numb,mess)

            elif 'bag' in name:
                numb = "8484037377"
                speak_take.speak(f"Whats the Message For {name}")
                mess = speak_take.takecommand()
                whatsapp.whatsapp(numb,mess)

            elif 'jarvis' in name:
                gro = "KsjFp8jhhPA9UHap9dRHcU"
                speak_take.speak(f"Whats the Message for {name} ")
                mess = speak_take.takecommand()
                whatsapp.Whatsapp_Grp(gro,mess)

TaskExe()


# Pong v3
#
# Pong is a game where two players try to bounce a moving ball into the oponents side
# The ball bounces off the walls and if it hits the oponents side the player gets a score
# The players use their paddles to stop and bounce the ball around
# First person to get to 11 points wins
#
# v3 requirements:
# In v3 the players are now able to move the paddles up and down by pressing the associated keys

# imports the pygame module
import pygame

# User-defined fucntions

def main():
   # initialize all pygame modules
   pygame.init()
   # create a pygame display window
   window_width = 500
   window_height = 400
   pygame.display.set_mode((window_width, window_height))
   # set the title of the display window
   pygame.display.set_caption('Pong')
   # get the display surface
   w_surface = pygame.display.get_surface()

   # create a game object
   game = Game(w_surface, window_width, window_height)
   # start the main game loop by calling the play method on the game object
   game.play()
   # quit pygame and clean up the pygame window
   pygame.quit()

# User-defined classes

class Game:
   # An object in this class represents a complete game.

   def __init__(self, surface, window_width, window_height):
      # Initialize a Game.
      # - self is the Game to initialize
      # - surface is the display window surface object
      # - bg_color us the background colour
      # - FPS is the frames per second
      # - game_Clock is our games timekeeping clock
      # - close_clicked is a boolean set as false
      # - self.continue_game is a boolean set as true

      # === general game objects
      self.surface = surface
      self.window_width = window_width
      self.window_height = window_height
      self.bg_color = pygame.Color('black')
      self.FPS = 60
      self.game_Clock = pygame.time.Clock()
      self.close_clicked = False
      self.continue_game = True

      # === game specific objects

      # Paddles specifications
      self.paddle_margin_wall = 100
      self.paddle_width = 12
      self.paddle_height = 50
      self.paddle_color = pygame.Color('white')

      # left paddle specifications
      self.left_paddle_position = [self.paddle_margin_wall, self.window_height //2 - self.paddle_height // 2]
      self.left_paddle_velocity = 0
      self.left_paddle_params = pygame.Rect(self.left_paddle_position[0], \
                                self.left_paddle_position[1], self.paddle_width, self.paddle_height)

      # right paddle specifications
      self.right_paddle_position = [self.window_width - self.paddle_width \
                                 - self.paddle_margin_wall, self.window_height //2 - self.paddle_height // 2]
      self.right_paddle_velocity = 0
      self.right_paddle_params = pygame.Rect(self.right_paddle_position[0], \
                                 self.right_paddle_position[1], self.paddle_width, self.paddle_height)

      # create the left and right paddles
      self.left_paddle = Paddle(self.paddle_color ,self.left_paddle_params, self.left_paddle_position, surface)
      self.right_paddle = Paddle(self.paddle_color ,self.right_paddle_params, self.right_paddle_position, surface)

      # balls specifications
      self.ball_radius = 6
      self.ball_colour = 'white'
      self.x_velocity = 9
      self.y_velocity = 3
      self.center = [window_width // 2 - 7, window_height // 2 - 7]

      # create the ball
      self.ball = Ball(self.ball_colour, self.ball_radius, self.center, \
                  [self.x_velocity, self.y_velocity], self.left_paddle, self.right_paddle, surface)

      # create scoreboard values
      self.score_left, self.score_right = 0, 0

   def play(self):
      # Play the game until the player presses the close box.
      # - self is the Game that should be continued or not.

      while not self.close_clicked:  # until player clicks close box
         # play frame
         self.handle_events()
         self.draw()
         if self.continue_game:
            self.update(self)
            self.decide_continue()
         self.game_Clock.tick(self.FPS) # run at most with FPS Frames Per Second

   def handle_events(self):
      # Handle each user event by changing the game state appropriately.
      # - self is the Game whose events will be handled
      # - pygame.KEYDOWN and pygame.KEYUP are functions used to check for keyboard presses

      events = pygame.event.get()
      for event in events:
         if event.type == pygame.QUIT:
            self.close_clicked = True

         if event.type == pygame.KEYDOWN and self.continue_game: # Has there been a keyboard click?
            self.handle_key_down(event.key)

         if event.type == pygame.KEYUP and self.continue_game: # Has there been a keyboard release?
            self.handle_key_up(event.key)

   def handle_key_down(self,key):
      # change the velocity of the paddles which moves the paddles up and down
      # depending on what key is pressed
      # velocity_change is used for this and is of type int
       # self.left_paddle_velocity and self.right_paddle_velocity are of type int

      velocity_change = 5

      # when the 'q' key is pressed down the left paddle moves up
      # when the 'a' key is pressed down the left paddle moves down
      if key == pygame.K_q:
         self.left_paddle_velocity = -velocity_change
      if key == pygame.K_a:
         self.left_paddle_velocity = velocity_change

      # when the 'p' key is pressed down the right paddle moves up
      # when the 'l' key is pressed down the right paddle moves down
      if key == pygame.K_p:
         self.right_paddle_velocity = -velocity_change
      if key == pygame.K_l:
         self.right_paddle_velocity = velocity_change

   def handle_key_up(self,key):
      # when the 'q' key or the 'a' is released the left paddle stops moving because the velocity is
      # changed to zero
      # self.left_paddle_velocity and self.right_paddle_velocity are of type int
      if key == pygame.K_q:
         self.left_paddle_velocity = 0
      if key == pygame.K_a:
         self.left_paddle_velocity = 0

      # when the 'p' key or the 'l' is released the right paddle stops moving
      if key == pygame.K_p:
         self.right_paddle_velocity = 0
      if key == pygame.K_l:
         self.right_paddle_velocity = 0

   def move_paddle(self):
      # sets the paddles to constantly move based on the velocities of the paddles and restricts
      # movement of paddles offscreen
      self.left_paddle_params.top +=  self.left_paddle_velocity
      self.right_paddle_params.top += self.right_paddle_velocity

      # restrict the left paddle from moving off screen
      if self.left_paddle_params.top <= 0:
         self.left_paddle_params.top = 0
      if self.left_paddle_params.bottom >= self.window_height:
         self.left_paddle_params.bottom = self.window_height

      # restrict the right paddle from moving off screen
      if self.right_paddle_params.top <= 0:
         self.right_paddle_params.top = 0
      if self.right_paddle_params.bottom >= self.window_height:
         self.right_paddle_params.bottom = self.window_height

   def draw(self):
      # Draw all game objects.
      # - self is the Game to draw
      self.surface.fill(self.bg_color) # clear the display surface first
      self.ball.draw()
      self.left_paddle.draw()
      self.right_paddle.draw()
      self.draw_score()
      pygame.display.update() # make the updated surface appear on the display

   def draw_score(self):
      # render scoreboards text to screen
      text_color = pygame.Color('white')
      text_font = pygame.font.SysFont("", 72)
      left_scoreboard_position = (0, 0)
      left_scoreboard = text_font.render(str(self.score_left), True, text_color)
      self.surface.blit(left_scoreboard , left_scoreboard_position)
      right_scoreboard = text_font.render(str(self.score_right), True, text_color)
      right_scoreboard_position = (self.window_width - right_scoreboard.get_width(), 0)
      self.surface.blit(right_scoreboard, right_scoreboard_position)

   def update(self, key):
      # Update the game objects for the next frame.
      # - self is the Game to update
      self.ball.move()
      self.move_paddle()
      self.ball.collision_paddle(self.left_paddle_params,self.right_paddle_params)

      # check if the ball hit the left or right wall and discribute points accordingly
      if self.center[0] < self.ball_radius:
         self.score_right +=  1
      if self.center[0] + self.ball_radius > self.window_width:
         self.score_left += 1

   def decide_continue(self):
      # Decides game end condition
      if self.score_left >= 11 or self.score_right >= 11:
         self.continue_game = False

class Ball:
   # An object in this class represents a Ball that that moves

   def __init__(self, ball_color, ball_radius, ball_center, ball_velocity, left_paddle, right_paddle, surface):
      # Initialize the ball.
      # - self is the Ball to initialize
      # - color is the pygame.Color of the Ball
      # - center is a list containing the x and y int
      #   coords of the center of the Ball
      # - radius is the int pixel radius of the Ball
      # - velocity is a list containing the x and y components
      # - surface is the window's pygame.Surface object
      self.color = pygame.Color(ball_color)
      self.radius = ball_radius
      self.center = ball_center
      self.velocity = ball_velocity
      self.left_padle = left_paddle
      self.right_paddle = right_paddle
      self.surface = surface

   def move(self):
      # Change the location of the Ball by adding the corresponding
      # speed values to the x and y coordinate of its center
      # - self is the Ball
      # move our Ball based on its velocity. If it comes in contact
      # with any of the four walls, reverse its direction of motion.
      size = self.surface.get_size()
      for index in range(0, len(size)):
         self.center[index] = (self.center[index] + self.velocity[index])
         if (self.center[index] < self.radius \
             or self.center[index] > size[index] - self.radius):
            self.velocity[index] = -self.velocity[index]

   def draw(self):
      # Draw the Ball on the surface
      # - self is the Ball
      pygame.draw.circle(self.surface, self.color, self.center, self.radius)

   def collision_paddle(self, left_paddle, right_paddle):
      # check if ball bounced on the  front of each paddle
      # collidepoint is associated with the pygame rect object and is used to chick if the ball
      # collides with the paddle
      if left_paddle.collidepoint(self.center[0] - self.radius, self.center[1]) and self.velocity[0] < 0 \
         or right_paddle.collidepoint(self.center[0] + self.radius, self.center[1]) and self.velocity[0] > 0: \
         self.velocity[0] = -self.velocity[0]

class Paddle:
   # An object in this class represents a paddle

   def __init__(self, paddle_color, paddle_params, paddle_position, surface):
      # Initialize the paddle
      # - self is the paddle to initialize
      # - paddle_color is the Color of the paddle
      # - rect_params give the values of rect_left, rect_top, rect_width, rect_height
      # - surface is the window's pygame.Surface object
      self.paddle_color = paddle_color
      self.paddle_params = paddle_params
      self.paddle_position = paddle_position
      self.surface = surface

   def draw(self):
      # Draw the paddle on the surface
      # - self is the paddle
      pygame.draw.rect(self.surface, self.paddle_color, self.paddle_params)

main()


import matplotlib.pyplot as plt
import itertools
import pandas as pd
import re
from scipy.stats import chi2_contingency
from statsmodels.sandbox.stats.multicomp import multipletests
plt.rcParams["font.family"] = "Lucida Sans Unicode"

class ProtocolData:
    def __init__(self, text, datetime_date):
        self.date = datetime_date
        self.member_affiliations = {}
        self.text = text
        self.interjection_locations = {}

        self.protocol_segments = []

        # Bit tricky
        self.words_preceding_interjections = {}

    def add_member(self, name, party):
        self.member_affiliations[name] = party

    def append_protocol_segment(self, segment):
        self.protocol_segments.append(segment)

    def count_word_preceding_interjection(self, word):
        self.words_preceding_interjections[word] = self.words_preceding_interjections.get(word, 0) + 1

    def print_data(self):
        print(self.date)
        print(self.member_affiliations)
        print(self.text)
        print(self.interjection_locations)
        print(self.words_preceding_interjections)
        print(len(self.interjection_locations))

    def print_data_graph(self):
        #drawnDict = dict(reversed(sorted(self.words_preceding_interjections.items(), key=lambda item: item[1])))
        #drawnDict = {k: v for k, v in drawnDict.items() if len(k) > 4 or k.isupper()}
        #drawnDict = dict(itertools.islice(drawnDict.items(), 15))
        #prediction = self.model.predict_output_word(['COMMENT'], topn = 20)

        #Pattern: Word, comment occurences, no comment occurences
        frequency_data = [(key, value, self.total_searched_words[key] - value) for (key, value) in self.most_frequent_words_near_search_words.items() if self.total_searched_words[key] > 3]
        df = pd.DataFrame(frequency_data, columns=["word", "comments", "no_comments"])
        df = df.set_index("word")
        significant_words = chisq_and_posthoc_corrected(df)
        significance_filtered_words = [(key, value) for (key, value) in self.most_frequent_words_near_search_words.items() if key in significant_words]
        relative_frequency_filtered_words = [(key, value / self.total_searched_words[key]) for (key, value) in significance_filtered_words if self.total_searched_words[key] > 3]
        #prediction = list(itertools.islice(self.most_frequent_words_near_search_words.items(), 20))
        prediction = sorted(relative_frequency_filtered_words[:20], key=lambda x: x[1], reverse=False)
        plt.barh([t[0] for t in prediction], [t[1] for t in prediction])
        #axes = plt.gca()
        #axes.set_ylim([0.011363, 0.011366])
        plt.show()

    def plot_significant_words(self):
        frequency_data = [(key, value, self.total_searched_words[key] - value) for (key, value) in self.most_frequent_words_near_search_words.items()]
        df = pd.DataFrame(frequency_data, columns=["word", "comments", "no_comments"])
        df = df.set_index("word")
        significant_words = chisq_and_posthoc_corrected(df)
        significance_filtered_words = [(key, value) for (key, value) in self.most_frequent_words_near_search_words.items() if key in significant_words if self.total_searched_words[key] > 3]
        relative_frequency_filtered_words = [(key, value / self.total_searched_words[key]) for (key, value) in significance_filtered_words]
        prediction = sorted(relative_frequency_filtered_words[:20], key=lambda x: x[1], reverse=False)
        plt.barh([t[0] for t in prediction], [t[1] for t in prediction])
        plt.show()

    def plot_word_count(self, search_pattern):
        pattern = re.compile(search_pattern)
        plotted_words = []
        for word, count in self.total_searched_words.items():
            if pattern.search(word):
                plotted_words.append([word, count])
        plt.barh([t[0] for t in plotted_words], [t[1] for t in plotted_words])
        plt.show()

    def get_full_text(self):
        full_text = ""
        for segment in self.protocol_segments:
            full_text += segment.get_text_equivalent()
        return full_text

def get_asterisks_for_pval(p_val):
    """Receives the p-value and returns asterisks string."""
    if p_val > 0.05:
        p_text = "ns"  # above threshold => not significant
    elif p_val < 1e-4:
        p_text = '****'
    elif p_val < 1e-3:
        p_text = '***'
    elif p_val < 1e-2:
        p_text = '**'
    else:
        p_text = '*'

    return p_text

# https://neuhofmo.github.io/chi-square-and-post-hoc-in-python/
def chisq_and_posthoc_corrected(df):
    """Receives a dataframe and performs chi2 test and then post hoc.
    Prints the p-values and corrected p-values (after FDR correction)"""
    # start by running chi2 test on the matrix
    chi2, p, dof, ex = chi2_contingency(df, correction=True)
    print(f"Chi2 result of the contingency table: {chi2}, p-value: {p}")

    avg_comments = df["comments"].mean()
    avg_no_comments = df["no_comments"].mean()
    # Any random occurence word, interjection afterwards?
    df.loc["AVERAGE_WORD"] = [avg_comments, avg_no_comments]
    print(df)

    #Bonferroni correction
    # post-hoc
    #all_combinations = list(itertools.com(df.index, 2))  # gathering all combinations for post-hoc chi2
    all_combinations = [(index, "AVERAGE_WORD") for index in df.index]
    p_vals = []
    print("Significance results:")
    for comb in all_combinations:
        new_df = df[(df.index == comb[0]) | (df.index == comb[1])]
        chi2, p, dof, ex = chi2_contingency(new_df, correction=True)
        p_vals.append(p)
        # print(f"For {comb}: {p}")  # uncorrected

    # checking significance
    # correction for multiple testing
    reject_list, corrected_p_vals = multipletests(p_vals, method='fdr_bh', alpha=0.05)[:2]
    for p_val, corr_p_val, reject, comb in zip(p_vals, corrected_p_vals, reject_list, all_combinations):
        print(f"{comb}: p_value: {p_val:5f}; corrected: {corr_p_val:5f} ({get_asterisks_for_pval(p_val)}) reject: {reject}")

    result = [comb[0] for p_val, corr_p_val, reject, comb in zip(p_vals, corrected_p_vals, reject_list, all_combinations) if reject]
    return result