import logging
from collections import Iterable
from typing import Optional, Union

from tqdm import tqdm
import pdb

from brainscore.metrics import Score
from brainscore.model_interface import BrainModel
from brainscore.utils import fullname
from model_tools.activations.pca import LayerPCA
from model_tools.brain_transformation import TemporalIgnore
from result_caching import store_xarray, store
import json
from model_tools.brain_transformation import LayerMappedModel

class LayerModel(BrainModel):
    def __init__(self, identifier, activations_model, layers):
        self.identifier = identifier
        self.activations_model = activations_model
        self.layers = layers

    def look_at(self, stimuli):
        activations = self.activations_model(stimuli, layers=self.layers)
        return activations

class ParamScores:
    def __init__(self, model_identifier, activations_model):
        self.model_identifier = model_identifier
        self._activations_model = activations_model
        self._logger = logging.getLogger(fullname(self))

    def __call__(self, benchmark_builder, params, layer, prerun=False):
        # Prerun is not supported as different benchmarks can have different inputs (although unlikely)
        assert not prerun, "Prerun not supported for ParamScores"
        benchmark_identifier = benchmark_builder().identifier
        return self._call(model_identifier=self.model_identifier,
                          model=self._activations_model, layer=layer,
                          benchmark_identifier=benchmark_identifier,
                          benchmark_builder=benchmark_builder, params=params)

    def build_layer_model(
            self, identifier, model, benchmark, layer):
        layer_model = LayerMappedModel(
                identifier=identifier,
                activations_model=model,
                region_layer_map={benchmark.region: layer})
        layer_model = TemporalIgnore(layer_model)
        return layer_model

    @store_xarray(
            identifier_ignore=['model', 'benchmark_builder', 'params'],
            combine_fields={'params': 'param'})
    def _call(self, model_identifier, layer, benchmark_identifier, # storage fields
              model, benchmark_builder, params):
        param_scores = []
        for param_str in tqdm(params, desc="params"):
            bench_args, bench_kwargs = json.loads(param_str)
            benchmark = benchmark_builder(*bench_args, **bench_kwargs)
            layer_model = self.build_layer_model(
                    identifier=f"{model_identifier}-{layer}",
                    model=model, benchmark=benchmark, layer=layer,
                    )
            score = benchmark(layer_model)
            score = score.expand_dims('param')
            score['param'] = [param_str]
            param_scores.append(score)
        param_scores = Score.merge(*param_scores)
        param_scores = param_scores.sel(param=params)  # preserve layer ordering
        return param_scores

class RegressParamScores(ParamScores):
    def build_layer_model(
            self, identifier, model, benchmark, layer):
        layer_model = LayerModel(
                identifier,
                model,
                [layer])
        return layer_model

class LayerParamScores:
    def __init__(self, model_identifier, activations_model):
        self.model_identifier = model_identifier
        self._activations_model = activations_model
        self._logger = logging.getLogger(fullname(self))

    def __call__(self, benchmark_builder, layer_and_params, prerun=False):
        # Params should be a tuple of (args, kwargs)
        # Prerun is not supported as different benchmarks can have different inputs (although unlikely)
        assert not prerun, "Prerun not supported for LayerParamScores"

        benchmark_identifier = benchmark_builder().identifier
        return self._call(model_identifier=self.model_identifier,
                          model=self._activations_model,
                          benchmark_builder=benchmark_builder,
                          benchmark_identifier=benchmark_identifier,
                          layer_and_params=layer_and_params)

    def build_layer_model(
            self, identifier, model, benchmark, layer):
        layer_model = LayerMappedModel(
                identifier=identifier,
                activations_model=model,
                region_layer_map={benchmark.region: layer})
        layer_model = TemporalIgnore(layer_model)
        return layer_model

    @store_xarray(
            identifier_ignore=['model', 'benchmark_builder', 'layer_and_params'],
            combine_fields={'layers': 'layer', 'params': 'param'})
    def _call(self, model_identifier, benchmark_identifier,  # storage fields
              model, benchmark_builder, layer_and_params):
        all_scores = []
        all_layer_param_str = []
        for layer, param_str in tqdm(layer_and_params, desc="layers"):
            bench_args, bench_kwargs = json.loads(param_str)
            benchmark = benchmark_builder(*bench_args, **bench_kwargs)
            layer_model = self.build_layer_model(
                    identifier=f"{model_identifier}-{layer}-layer-param",
                    model=model, benchmark=benchmark, layer=layer,
                    )
            score = benchmark(layer_model)
            score = score.expand_dims('layer_param')
            layer_param_str = '%s-%s' % (layer, param_str)
            score['layer_param'] = [layer_param_str]
            all_scores.append(score)
            all_layer_param_str.append(layer_param_str)
        all_scores = Score.merge(*all_scores)
        all_scores = all_scores.sel(layer_param=all_layer_param_str)
        return all_scores

class LayerRegressParamScores(LayerParamScores):
    def build_layer_model(
            self, identifier, model, benchmark, layer):
        layer_model = LayerModel(
                identifier,
                model,
                [layer])
        return layer_model

class LayerActivations:
    def __init__(self, model_identifier, activations_model):
        self.model_identifier = model_identifier
        self._activations_model = activations_model
        self._logger = logging.getLogger(fullname(self))

    def __call__(
            self, benchmark, layers, benchmark_identifier=None, prerun=False):
        return self._call(
                model_identifier=self.model_identifier,
                benchmark_identifier=benchmark_identifier or benchmark.identifier,
                model=self._activations_model,
                benchmark=benchmark, layers=layers, prerun=prerun)

    @store_xarray(
            identifier_ignore=['model', 'benchmark', 'layers', 'prerun'],
            combine_fields={'layers': 'layer'})
    def _call(self, model_identifier, benchmark_identifier,  # storage fields
              model, benchmark, layers, prerun=False):
        if prerun:
            # pre-run activations together to avoid running every layer separately
            model(layers=layers, stimuli=benchmark._assembly.stimulus_set)

        for layer in tqdm(layers, desc="layers"):
            layer_model = LayerMappedModel(
                    identifier=f"{model_identifier}-{layer}",
                    activations_model=model,
                    region_layer_map={benchmark.region: layer})
            layer_model = TemporalIgnore(layer_model)
            layer_model.start_recording(
                    benchmark.region,
                    time_bins=benchmark.timebins)
            layer_model.look_at(benchmark._assembly.stimulus_set)


import os
import csv
import imageio
import scipy.misc
import cv2 as cv
import util_functions
import numpy as np
import matplotlib.pyplot as plt

def down_sample(img, mult_num, dim=3):
    if dim == 3:
        h, w, c = img.shape
    else:
        h, w = img.shape
    new_h = h // mult_num
    new_w = w // mult_num
    img = scipy.misc.imresize(img, (new_h, new_w))
    return img

def read_polygon_info(file_name, city_name):
    with open(file_name, 'r') as f:
        reader = csv.reader(f)
        for line in reader:
            if city_name in line[0]:
                line_num = int(line[2])
                line_record = []
                for i in range(line_num - 1):
                    p_pair = list()
                    x = line[3 + i*2]
                    y = line[3 + i*2 + 1]
                    if str(x) != 'NaN' and str(y) != 'NaN':
                        p_pair.append(float(x))
                        p_pair.append(float(y))
                        line_record.append(p_pair)
                    else:
                        pass
                        line_record_temp = line_record
                        line_record = [line_record_temp[-1]]
                        print(line_record_temp)
                        yield line_record_temp
                #yield line_record

def add_polygons(img, reader):
    for coords in reader:
        vrx = np.array(coords, dtype='int32')
        # img = cv.fillPoly(img, [vrx], (1, 0, 0))
        img = cv.polylines(img, [vrx], False, (1, 0, 0), thickness=20)
    return img[:, :, 0]

if __name__ == '__main__':
    city_name = 'Arlington'
    img_dir = r'/home/lab/Documents/bohao/data/arlington'
    check_result = True
    for img_id in [1, 2, 3]:
        img_name = '{}_0{}.tif'.format(city_name, img_id)
        gt_file_name = '{}_0{}_roadCoord.csv'.format(city_name, img_id)

        info_reader = read_polygon_info(os.path.join(img_dir, gt_file_name), city_name)
        img = imageio.imread(os.path.join(img_dir, img_name))

        gt_img = np.zeros(img.shape, np.uint8)
        gt_img = add_polygons(gt_img, info_reader)

        img = down_sample(img, 2)
        gt_img = down_sample(gt_img, 2, dim=2)

        new_img_name = '{}{}_RGB.tif'.format(city_name, img_id)
        new_gt_name = '{}{}_GT.png'.format(city_name, img_id)

        imageio.imsave(os.path.join(img_dir, new_img_name), img)
        imageio.imsave(os.path.join(img_dir, new_gt_name), gt_img)

        if check_result:
            plt.figure(figsize=(10, 8))
            # ax1 = plt.subplot(211)
            img = util_functions.add_mask(img, gt_img, [255, 0, 0], mask_1=1)
            plt.imshow(img)
            plt.axis('off')
            '''ax2 = plt.subplot(212, sharex=ax1, sharey=ax1)
            plt.imshow(gt_img)
            plt.axis('off')'''
            plt.tight_layout()
            plt.show()


from PyQt5.QtGui import QPainter, QFont, QColor, QPixmap, QPen, QBrush
from PyQt5.QtCore import Qt, QRect, QPoint, QRectF

class Drawable:
    def draw(self, ctx, location, size):
        pass

class Text(Drawable):
    def __init__(self, text, font, color):
        super().__init__()
        self.text = text
        self.font = font
        self.color = color

    def draw(self, ctx, location, size=None):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(self.color)
        qp.setFont(self.font)
        qp.drawText(location[0], location[1], self.text)

class CenterText(Text):
    def draw(self, ctx, location, size=None):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(self.color)
        qp.setFont(self.font)
        qp.drawText(QRect(0, 0, 1600, 900), Qt.AlignCenter, self.text)

class CodeDrawable(Drawable):
    def __init__(self):
        super().__init__()
        self.code = ''
    def setCode(self, code):
        self.code = code
    def draw(self, ctx, location, size=None):
        nowx = location[0]
        i = 0
        event = ctx[0]
        qp = ctx[1]
        COND = QColor("#D59DF6")
        METHOD = QColor("#71B1FE")
        NORM = QColor("#FFFFFF")
        fontSize = 16
        lineSpacing = 1.5
        qp.setFont(QFont('D2Coding', fontSize))
        wSpacing = 0.65
        nowx = location[0]
        nowy = location[1]
        code = self.code
        while i<len(code):
            if code[i] == '\n':
                nowx = location[0]
                nowy += 1*lineSpacing*fontSize
                i+=1
            elif code[i:i+2] == 'if':
                qp.setPen(COND)
                qp.drawText(nowx, nowy, code[i:i+2])
                nowx += fontSize * 2 * wSpacing
                i += 2
            elif code[i:i+4] == 'else':
                qp.setPen(COND)
                qp.drawText(nowx, nowy, code[i:i+4])
                nowx += fontSize * 4 * wSpacing
                i += 4
            elif code[i:i+5] == 'while':
                qp.setPen(COND)
                qp.drawText(nowx, nowy, code[i:i+5])
                nowx += fontSize * 5 * wSpacing
                i += 5
            elif code[i:i+2] == 'in':
                qp.setPen(COND)
                qp.drawText(nowx, nowy, code[i:i+2])
                nowx += fontSize * 2 * wSpacing
                i += 2
            elif code[i:i+3] == 'for':
                qp.setPen(COND)
                qp.drawText(nowx, nowy, code[i:i+3])
                nowx += fontSize * 3 * wSpacing
                i += 3
            elif code[i:i+5] == 'range':
                qp.setPen(METHOD)
                qp.drawText(nowx, nowy, code[i:i+5])
                nowx += fontSize * 5 * wSpacing
                i += 5
            elif code[i:i+6] == 'attack':
                qp.setPen(METHOD)
                qp.drawText(nowx, nowy, code[i:i+6])
                nowx += fontSize * 6 * wSpacing
                i += 6
            elif code[i:i+7] == 'defence':
                qp.setPen(METHOD)
                qp.drawText(nowx, nowy, code[i:i+7])
                nowx += fontSize * 7 * wSpacing
                i += 7
            elif code[i:i+6] == 'player':
                qp.setPen(METHOD)
                qp.drawText(nowx, nowy, code[i:i+6])
                nowx += fontSize * 6 * wSpacing
                i += 6
            elif code[i:i+5] == 'enemy':
                qp.setPen(METHOD)
                qp.drawText(nowx, nowy, code[i:i+5])
                nowx += fontSize * 5 * wSpacing
                i += 5
            else:
                qp.setPen(NORM)
                qp.drawText(nowx, nowy, code[i:i+1])
                nowx += fontSize * 1 * wSpacing
                i += 1

class Image(Drawable):
    def __init__(self, url):
        super().__init__()
        self.pixmap = QPixmap(url)

    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.drawPixmap(QRect(location[0], location[1], size[0], size[1]), self.pixmap)

class AnimatedImage(Drawable):
    def __init__(self, urls):
        super().__init__()
        self.pixmap = []
        self.frame = 0
        for url in urls:
            self.pixmap.append(QPixmap(url))

    def draw(self, ctx, location, size, seeingRight = True):
        event = ctx[0]
        qp = ctx[1]
        if seeingRight:
            qp.drawPixmap(QRect(location[0], location[1], size[0], size[1]), self.pixmap[int(self.frame)])
        else:
            qp.drawPixmap(QRect(location[0], location[1], size[0], size[1]), QPixmap.fromImage(self.pixmap[int(self.frame)].toImage().mirrored(horizontal = True, vertical = False)))
    def animate(self, i):
        if self.frame+i >= len(self.pixmap):
            self.frame = (self.frame+i) % len(self.pixmap)
            return True
        else:
            self.frame = self.frame + i
            return False

    def reset(self):
        self.frame = 0

class BarrierDrawable(AnimatedImage):
    def __init__(self):
        super().__init__([f'./res/image/shield/pipo-btleffect207_192-{i}.png' for i in range(20)])
    def draw(self, ctx, location, size):
        ctx[1].setOpacity(0.45)
        super().draw(ctx, location, size)
        ctx[1].setOpacity(1)

class Rect(Drawable):
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setBrush(QColor(25, 0, 90, 200))
        qp.drawRect(location[0], location[1], size[0], size[1])

class ScrollDrawable(Drawable):
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200), 5))
        qp.setBrush(QColor(30, 30, 30))
        qp.drawRect(location[0], location[1], size[0], size[1])

class BlockDrawable(Drawable):
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200), 3))
        qp.setBrush(QColor(30, 30, 30))
        qp.drawRect(location[0], location[1], size[0], size[1])

    def drawWrap(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200, 50), 3))
        qp.setBrush(QColor(200, 200, 200, 50))
        qp.drawRect(location[0], location[1], size[0], size[1])

class TitleButtonDrawable(Drawable):
    def draw(self, ctx, location, size=(230, 100)):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(0, 0, 0, 150), 3))
        qp.setBrush(QColor(0, 0, 0, 150))
        qp.drawRect(location[0], location[1], size[0], size[1])

    def drawWrap(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200, 50), 3))
        qp.setBrush(QColor(200, 200, 200, 50))
        qp.drawRect(location[0], location[1], size[0], size[1])

class HpDrawable(Drawable):
    def __init__(self, width, height):
        self.width = width
        self.height = height
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(0, 0, 0), 2))
        qp.setBrush(QColor(20, 20, 20))
        qp.drawRect(location[0], location[1], self.width, self.height)
        qp.setPen(QPen(QColor(0, 0, 0), 2))
        qp.setBrush(QColor(255, 40, 40))
        qp.drawRect(location[0], location[1], size[0], self.height)

class ShopDrawable(Drawable):
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200), 5))
        qp.setBrush(QColor(30, 30, 30))
        qp.drawRect(location[0], location[1], size[0], size[1])

class ShopButtonDrawable(Drawable):
    def __init__(self, category=10):
        super().__init__()
        self.category = category
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200), 3))
        category = self.category
        if category == 10:
            color = QColor(30, 30, 30)
        elif category == 0:
            color = QColor(30, 30, 70)
        elif category == 1 or category == 2:
            color = QColor(30, 30, 120)
        elif category == 3:
            color = QColor(40, 40, 150)
        elif category == 11:
            color = QColor(80, 50, 30)
        else:
            color = QColor(100, 100, 100)
        qp.setBrush(color)
        qp.drawRect(location[0], location[1], size[0], size[1])

    def drawWrap(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200, 50), 3))
        qp.setBrush(QColor(200, 200, 200, 50))
        qp.drawRect(location[0], location[1], size[0], size[1])

    def drawSold(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(0, 0, 0, 200), 3))
        qp.setBrush(QColor(0, 0, 0, 50))
        qp.drawRect(location[0], location[1], size[0], size[1])

class ButtonDrawable(Drawable):
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(200, 200, 200), 3))
        qp.setBrush(QColor(30, 30, 30))
        qp.drawRect(location[0], location[1], size[0], size[1])

class ShopMoneyDrawable(Drawable):
    def draw(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(255, 255, 0), 10))
        qp.setBrush(QColor(185, 185, 0))
        qp.drawRect(location[0]+10, location[1]+10, size[0]-20, size[1]-20)

    def draw2(self, ctx, location, size):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(255, 255, 255), 2))
        qp.setBrush(QColor(0, 0, 0))
        qp.drawRect(location[0], location[1], size[0], size[1])

class CostDrawable(Drawable):
    def __init__(self, cost):
        self.cost = cost
    def draw(self, ctx, location, size=(40, 40)):
        event = ctx[0]
        qp = ctx[1]
        qp.setPen(QPen(QColor(255, 255, 0), 4))
        qp.setBrush(QColor(185, 185, 0))
        qp.drawRect(location[0]-size[0]//2, location[1]-size[1]//2, size[0], size[1])
        Text(str(self.cost), QFont('NotoMono', 24, 1), QColor(0, 0, 0)).draw(ctx, (location[0]-10, location[1]+12))


# -*- coding: utf-8 -*-
"""
Let's fit a Gaussian process without the explicit solution'
"""

import numpy as np
from scipy.stats import multivariate_normal
import matplotlib.pyplot as plt
import math

def target_func(x):
    omega = 1
    return np.sin(2*np.pi*omega*x)

def make_data(train_points,noise_param):
    return target_func(train_points)+np.random.normal(loc=0,scale=noise_param,size=len(train_points))

def kernel(x,y):
    eps = 10**(-12)
    return np.exp(-abs(x-y)**2/(2*0.2**2))+eps*(abs(x-y)<0.01)

def compute_prior_covariance(points):
    N = len(points)
    covar = np.zeros((N,N))
    for i in range(N):
        for j in range(N):
            covar[i][j] = kernel(points[i],points[j])
    return covar

def set_up_prior(test_points):
    mean = np.array([0]*len(test_points))
    covar = np.array(compute_prior_covariance(test_points))
    return mean,covar

def matching_squared_sum(u,y,mesh_size):
    total = 0
    for x in y.keys():
        total += (y[x]-u[int(x/mesh_size)])**2
    return total

def unnormalized_posterior(u,y,prior_mean,prior_covar,noise_param,mesh_size):
    prior_eval = multivariate_normal.pdf(u,mean = prior_mean,cov=prior_covar)
    conditional_eval = np.exp((-len(y.keys())/(2*noise_param**2))*matching_squared_sum(u, y, mesh_size))
    return prior_eval*conditional_eval

#def gaussian_kernel_sample(x):
 #   z = np.random.normal(loc = x,scale = 0.1)
  #  return z

def crank_nicholson(x,prior_covar):
    beta = 0.12
    return np.sqrt(1-beta**2)*x+beta*np.random.multivariate_normal([0]*len(x),prior_covar)

def rho(x,z,y,prior_mean,prior_covar,noise_param,mesh_size):
    #z is the new proposal, x is the current state, y is the data
    proposal_score = np.exp((1/(2*noise_param**2))*(matching_squared_sum(x, y,mesh_size)-matching_squared_sum(z, y, mesh_size)))
    return min(proposal_score,1)

def metropolis_hastings_sample(N,y,prior_mean,prior_covar,noise_param,mesh_size,mh_kernel):
    accepted = 0
    rejected = 0
    test_points = np.linspace(0,1,21)
    x = prior_mean
    samples = []
    i = 0
    while i < N:
        z = mh_kernel(x,prior_covar)
        if rho(x,z,y,prior_mean,prior_covar,noise_param,mesh_size) > np.random.random():
            x = z
            samples.append(x)
            accepted += 1
        else:
            samples.append(x)
            rejected += 1
        i += 1
    return samples

def gaussian_kernel(x,y):
    eps = 0 #10**(-11)
    return 0.5*np.exp(-(x-y)**2/(2*0.3**2))+eps*(abs(x-y)<0.01)

def analytic_fitting(data,test_points,noise_param):
    def K(A,B):
        kernel_matrix = list(np.zeros((len(A),len(B))))
        for i in range(len(A)):
            for j in range(len(B)):
                kernel_matrix[i][j] = kernel(A[i],B[j])
        return np.array(kernel_matrix)
    y = [data[key] for key in data.keys()]
    noise_matrix = noise_param**2*np.identity(6)
    X = np.linspace(0,1,6)
    return K(test_points,X)@np.linalg.inv((K(X,X)+noise_matrix))@y

def mse(x,y):
    return np.sum((x-y)**2)/len(x)

def GR(chains):
    N = chains.shape[1]
    chain_vars = np.var(chains,axis=1,ddof=1)
    avg_vars = np.mean(chain_vars,axis=0)
    chain_means = np.mean(chains,axis=1)
    B_n = np.var(chain_means,axis=0,ddof=1)
    sig_sqr = ((N-1)/N)*avg_vars+B_n
    R = np.sqrt(np.divide(sig_sqr,avg_vars))
    return np.mean(R)

def main():

    mesh_size = 0.05
    train_point_count = 6
    noise_param = 0.25
    train_points = np.linspace(0,1,train_point_count)
    y =  dict(zip(train_points,make_data(train_points, noise_param)))
    samples_per_chain = 2000
    burn_in = 1000
    n_chains = 10

    test_point_count = int(1/mesh_size+1)
    test_points = np.linspace(0,1,test_point_count)
    prior_mean, prior_covar = set_up_prior(test_points)

    chains = []
    for i in range(n_chains):
        chain=metropolis_hastings_sample(samples_per_chain,y,prior_mean,prior_covar,noise_param,mesh_size,crank_nicholson)[burn_in:]
        chains.append(chain)
    chains = np.array(chains)

    y_keys = y.keys()
    true_y = [target_func(x) for x in test_points]
    plt.scatter(test_points,true_y,label = "True function, y = sin(2\u03C0x)",marker="o",s=20)
    plt.scatter(y_keys,[y[key] for key in y_keys],label="Data",marker="x",s=20)

    #plt.scatter(test_points,prior_mean)
    #print(chains.shape)
    posterior_experimental_mean = np.mean(chains,axis=(0,1))
    #print(posterior_experimental_mean)

    #plt.scatter(test_points,posterior_samples[-1],label="Posterior sample")
    plt.scatter(test_points,posterior_experimental_mean,label="Posterior MH mean",marker="s",s=20)

    analytic_mean = analytic_fitting(y,test_points,noise_param)
    #print(np.linalg.norm(posterior_experimental_mean-analytic_mean))
    plt.scatter(test_points,analytic_mean,label="Analytic mean",marker="*",s=20)
    plt.legend()
    plt.xlabel("x")
    plt.ylabel("y")
    plt.title("Metropolis-Hastings Mean GP")

    print(GR(chains))
    #analytic_sample_errors = [mse(analytic_mean,posterior_sample) for posterior_sample in posterior_samples]

    #plt.scatter([i for i in range(len(analytic_sample_errors))],analytic_sample_errors)

main()


from flask import Flask, request, abort, render_template, jsonify

import os
import json
import hashlib
import hmac

import git

import openai
from openai.error import RateLimitError

app = Flask(__name__)
w_secret = str(os.environ.get('WEBHOOK_SECRET'))
openai.api_key = str(os.getenv("OPENAI_API_KEY"))

@app.route("/")
def hello_world():
    return "
henLO world! gpt api
"

# url_for('static', filename='style.css')

@app.route("/pong")
def pong():
    return render_template('pong.html')

@app.route("/henlo")
def henlo():
    return "
henLO
"

@app.route('/gpt4', methods=['GET', 'POST'])
def gpt4():
    user_input = request.args.get(
        'user_input') if request.method == 'GET' else request.form['user_input']
    messages = [{"role": "user", "content": user_input}]

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=messages
        )
        content = response.choices[0].message["content"]
    except RateLimitError:
        content = "The server is experiencing a high volume of requests. Please try again later."

    return jsonify(content=content)

def is_valid_signature(x_hub_signature, data, private_key):
    hash_algorithm, github_signature = x_hub_signature.split('=', 1)
    algorithm = hashlib.__dict__.get(hash_algorithm)
    encoded_key = bytes(private_key, 'latin-1')
    mac = hmac.new(encoded_key, msg=data, digestmod=algorithm)
    return hmac.compare_digest(mac.hexdigest(), github_signature)

@app.route('/update_server', methods=['POST'])
def webhook():
    if request.method != 'POST':
        return 'OK'
    else:
        abort_code = 418
        # Do initial validations on required headers
        if 'X-Github-Event' not in request.headers:
            abort(abort_code)
        if 'X-Github-Delivery' not in request.headers:
            abort(abort_code)
        if 'X-Hub-Signature' not in request.headers:
            abort(abort_code)
        if not request.is_json:
            abort(abort_code)
        if 'User-Agent' not in request.headers:
            abort(abort_code)
        ua = request.headers.get('User-Agent')
        if not ua.startswith('GitHub-Hookshot/'):
            abort(abort_code)

        event = request.headers.get('X-GitHub-Event')
        if event == "ping":
            return json.dumps({'msg': 'Hi!'})
        if event != "push":
            return json.dumps({'msg': "Wrong event type"})

        x_hub_signature = request.headers.get('X-Hub-Signature')
        # webhook content type should be application/json for request.data to have the payload
        # request.data is empty in case of x-www-form-urlencoded
        if not is_valid_signature(x_hub_signature, request.data, w_secret):
            print('Deploy signature failed: {sig}'.format(sig=x_hub_signature))
            abort(abort_code)

        payload = request.get_json()
        if payload is None:
            print('Deploy payload is empty: {payload}'.format(
                payload=payload))
            abort(abort_code)

        # if payload['ref'] != 'refs/heads/master':
        #    return json.dumps({'msg': 'Not master; ignoring'})

        repo = git.Repo('~/pythonanywhere')
        origin = repo.remotes.origin

        pull_info = origin.pull()

        if len(pull_info) == 0:
            return json.dumps({'msg': "Didn't pull any information from remote!"})
        if pull_info[0].flags > 128:
            return json.dumps({'msg': "Didn't pull any information from remote!"})

        commit_hash = pull_info[0].commit.hexsha
        build_commit = f'build_commit = "{commit_hash}"'
        print(f'{build_commit}')
        return 'Updated PythonAnywhere server to commit {commit}'.format(commit=commit_hash)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000, debug=True)


from collections import Counter

def frequency_sort(items):
    counts = Counter(items)
    return sorted(items, key=lambda k: counts[k] * len(items) - items.index(k), reverse=True)

if __name__ == '__main__':
    print("Example:")
    print(frequency_sort([4,6,2,2,2,6,4,4,4]))


class A:
    message="Class Message"

    @classmethod
    def cfoo(cls,):
        print(cls.message)

    def foo(self, msg):
        self.message=msg
        print(self.message)

    def __str__(self):
        return self.message

if __name__ == "__main__":

    a = A()
    print(a)  # Class Message

    print(a.cfoo())

    print(A.cfoo())


import dicom2nifti
import os
import gzip
import shutil

fromPath = 'F:/Research/data/raw'

for x in range(2, 3):
    path = fromPath + '/' + str(x) + '/CT/separate'
    toPath = fromPath + '/' + str(x) + '/CT/rawnii/'
    if os.path.isdir(path):
        firstFolder = os.listdir(path)
        singleFolderPath = path + '/' + firstFolder[0] + '/'
        if not os.path.isdir(toPath):
            os.mkdir(toPath)
        dicom2nifti.convert_directory(singleFolderPath, toPath)

        rawnii = os.listdir(toPath)
        if len(rawnii) > 1:
            print('Error, more than 1 file output ' + str(x))

        fileName = toPath + '/' + rawnii[0]
        newFileName = toPath + '/raw.nii'

        with gzip.open(fileName, 'rb') as f_in:
            with open(newFileName, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        os.remove(fileName)

            #dicom2nifti.dicom_series_to_nifti(singleFolderPath, toPath, reorient_nifti=True)
        #except:
            #print("An exception occurred: " + str(x))


class Solution:
    def oddString(self, words: List[str]) -> str:
        diff = defaultdict(list)

        for word in words:
            d = []

            for i in range(1, len(word)):
                d.append(ord(word[i]) - ord(word[i - 1]))

            diff[tuple(d)].append(word)

        for k, v in diff.items():
            if len(v) == 1:
                return v[0]


# https://leetcode.com/problems/range-sum-query-immutable/

### Description:
# Given an integer array nums, handle multiple queries of the following type:
# Calculate the sum of the elements of nums between indices left and right inclusive where left <= right.
# Implement the NumArray class:
# NumArray(int[] nums) Initializes the object with the integer array nums.
# int sumRange(int left, int right)
# Returns the sum of the elements of nums between indices left and right inclusive (i.e. nums[left] + nums[left + 1] + ... + nums[right]).

### Examples:
# Input
# ["NumArray", "sumRange", "sumRange", "sumRange"]
# [[[-2, 0, 3, -5, 2, -1]], [0, 2], [2, 5], [0, 5]]
# Output
# [null, 1, -1, -3]

# Explanation
# NumArray numArray = new NumArray([-2, 0, 3, -5, 2, -1]);
# numArray.sumRange(0, 2); // return (-2) + 0 + 3 = 1
# numArray.sumRange(2, 5); // return 3 + (-5) + 2 + (-1) = -1
# numArray.sumRange(0, 5); // return (-2) + 0 + 3 + (-5) + 2 + (-1) = -3

### Solution:
# https://leetcode.com/problems/range-sum-query-immutable/discuss/75184/5-lines-C%2B%2B-4-lines-Python

# Create an array acc that stores the accumulated sum for nums such that:
# acc[i] = nums[0] + ... + nums[i-1]
# Then just return acc[j+1] - acc[i] in sumRange

class NumArray(object):
    def __init__(self, nums):
        self.acc = [0]
        for num in nums:
            self.acc.append(self.acc[-1] + num)

    def sumRange(self, i, j):
        return self.acc[j + 1] - self.acc[i]

numArray = NumArray(nums=[-2, 0, 3, -5, 2, -1])
print(numArray.sumRange(0, 2)) # 1
print(numArray.sumRange(2, 5)) # -1
print(numArray.sumRange(0, 5)) # -3


from conan import ConanFile
from conan.tools.microsoft import check_min_vs, is_msvc_static_runtime, is_msvc
from conan.tools.files import get, copy, rm, rmdir
from conan.tools.build import check_min_cppstd
from conan.tools.cmake import CMake, CMakeDeps, CMakeToolchain, cmake_layout
from conan.tools.files import apply_conandata_patches, export_conandata_patches
import os

required_conan_version = ">=1.53.0"

class CpptraceConan(ConanFile):
    name = "cpptrace"
    description = "Simple, portable, and self-contained stacktrace library for C++11 and newer "
    license = ("MIT", "LGPL-2.1-only", "BSD-2-Clause-Views")
    url = "https://github.com/conan-io/conan-center-index"
    homepage = "https://github.com/jeremy-rifkin/cpptrace"
    topics = ("stacktrace", "backtrace", "stack-trace", "back-trace", "trace", "utilities", "error-handling")
    package_type = "library"
    settings = "os", "arch", "compiler", "build_type"
    options = {
        "shared": [True, False],
        "fPIC": [True, False],
    }
    default_options = {
        "shared": False,
        "fPIC": True,
    }

    @property
    def _min_cppstd(self):
        return 11

    def config_options(self):
        if self.settings.os == "Windows":
            del self.options.fPIC

    def configure(self):
        if self.options.shared:
            self.options.rm_safe("fPIC")

    def layout(self):
        cmake_layout(self, src_folder="src")

    def requirements(self):
        self.requires("libdwarf/0.8.0")

    def validate(self):
        if self.settings.compiler.cppstd:
            check_min_cppstd(self, self._min_cppstd)
        check_min_vs(self, 191)

    def source(self):
        get(self, **self.conan_data["sources"][self.version], strip_root=True)

    def export_sources(self):
        export_conandata_patches(self)

    def generate(self):
        tc = CMakeToolchain(self)
        if is_msvc(self):
            tc.variables["USE_MSVC_RUNTIME_LIBRARY_DLL"] = not is_msvc_static_runtime(self)
        if not self.options.shared:
            tc.variables["CPPTRACE_STATIC"] = True
        tc.variables["CPPTRACE_USE_SYSTEM_LIBDWARF"] = True
        tc.generate()
        tc = CMakeDeps(self)
        tc.generate()

    def build(self):
        apply_conandata_patches(self)
        cmake = CMake(self)
        cmake.configure()
        cmake.build()

    def package(self):
        copy(self, pattern="LICENSE", dst=os.path.join(self.package_folder, "licenses"), src=self.source_folder)
        cmake = CMake(self)
        cmake.install()

        if self.settings.os == "Windows" and self.options.shared:
            copy(self, "*.dll", src=self.build_folder, dst=os.path.join(self.package_folder, "bin"), keep_path=False)

        rmdir(self, os.path.join(self.package_folder, "lib", "pkgconfig"))
        rmdir(self, os.path.join(self.package_folder, "lib", "cmake"))
        rmdir(self, os.path.join(self.package_folder, "share"))
        rm(self, "*.la", os.path.join(self.package_folder, "lib"))
        rm(self, "*.pdb", os.path.join(self.package_folder, "lib"))
        rm(self, "*.pdb", os.path.join(self.package_folder, "bin"))

    def package_info(self):
        self.cpp_info.libs = ["cpptrace"]

        self.cpp_info.set_property("cmake_module_file_name", "cpptrace")
        self.cpp_info.set_property("cmake_module_target_name", "cpptrace")

        if self.settings.os in ["Linux", "FreeBSD"]:
            self.cpp_info.system_libs.append("m")
            self.cpp_info.system_libs.append("pthread")
            self.cpp_info.system_libs.append("dl")
        if self.settings.os == "Windows":
            self.cpp_info.system_libs.append("dbghelp")

        # TODO: to remove in conan v2 once cmake_find_package_* generators removed
        self.cpp_info.filenames["cmake_find_package"] = "CPPTRACE"
        self.cpp_info.filenames["cmake_find_package_multi"] = "cpptrace"
        self.cpp_info.names["cmake_find_package"] = "CPPTRACE"
        self.cpp_info.names["cmake_find_package_multi"] = "cpptrace"


# Databricks notebook source
# MAGIC %run ./sampling/cdf_sampling

# COMMAND ----------

from functools import partial, reduce
from typing import Iterable, Callable, List

from pyspark.sql import DataFrame, functions as F

# Prevent spark from auto broadcasting the metadata
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)

# COMMAND ----------

union_dfs = partial(reduce, DataFrame.unionByName)

def get_pivot_field_values(df: DataFrame, pivot_field: str) -> List[str]:
  """Get a list of the distinct values in a field to be used in the values argument of a
  call to DataFrame.pivot. These values will be the columns in the pivoted dataframe.

  Args:
      df (pyspark.sql.DataFrame): DataFrame to be pivoted
      pivot_field (str): Field to pivot by. Distinct values will be the columns of the pivoted dataframe

  Returns:
      List[str]: Distinct values in the pivot field
  """
  pivot_field_values = (
    df
    .select(pivot_field)
    .distinct()
    .collect()
  )
  pivot_field_values = [x[pivot_field] for x in pivot_field_values]
  return pivot_field_values

def generate_empty_samples(
  metadata_df,
  sample_id_df,
  variable_cols=("TABLE_NAME", "FIELD_NAME"),
) -> DataFrame:
  """ Cross-joins the variable names in the metadata with the rows of sample_id_df,
  so that each row in sample_id_df is associated with each distinct combination
  of variables in the metadata.

  This doesnt actually do the sampling, it simply builds a tidy dataframe that
  the samples are generated onto (see the example).

  Args:
      metadata_df (pyspark.sql.DataFrame): Contains metadata
      sample_id_df (pyspark.sql.DataFrame): Contains indexes to generate samples for

  Returns:
      pyspark.sql.DataFrame:  Cariable names in the metadata cross-joined with the rows of sample_id_df

  Example:
      >> metadata_df
      | TABLE_NAME | FIELD_NAME | ... |
      | ---------- | ---------- | --- |
      |  table1    |   field1   | ... |
      |  table1    |   field2   | ... |
      |  table2    |   field1   | ... |

      >> sample_id_df
      | SAMPLE_ID |
      | --------- |
      |     1     |
      |     2     |

      >> generate_empty_samples(metadata_df, sample_id_df)
      | SAMPLE_ID | TABLE_NAME | FIELD_NAME |
      | --------- | ---------- | ---------- |
      |     1     |  table1    |   field1   |
      |     1     |  table1    |   field2   |
      |     1     |  table2    |   field1   |
      |     2     |  table1    |   field1   |
      |     2     |  table1    |   field2   |
      |     2     |  table2    |   field1   |

  """
  variable_df = (
    metadata_df
    .select(*variable_cols)
    .distinct()
  )

  sample_df = sample_id_df.crossJoin(variable_df.hint("broadcast"))

  return sample_df

def continuous_field_generator(meta_df: DataFrame, base_df: DataFrame, variable_fields: Iterable[str]) -> DataFrame:
  """Generate a dataframe of (long-formatted) values for fields with VALUE_TYPE='CONTINUOUS' in the metadata.
  This is done by uniform sampling from bins defining the percentiles of a continuous distribution (i.e.
  first we pick a bin with 1% width, then we sample a value from within that bin).

  Args:
      meta_df (pyspark.sql.DataFrame): Metadata for the dataset
      base_df (pyspark.sql.DataFrame): Base data to sample values for, indexed over patients / events.
      variable_fields (Iterable[str]): Names of metadata fields that define the fields to samples (e.g.
        we generally want to sample per combination of 'TABLE_NAME' and 'FIELD_NAME')

  Returns:
      DataFrame: DataFrame of generated data (long-formatted), with a single value per
        unique combination of variable_fields.
  """
  filtered_meta_df = meta_df.filter(F.col("VALUE_TYPE") == "CONTINUOUS")
  sample_df = generate_empty_samples(filtered_meta_df, base_df, variable_cols=variable_fields)

  # Sampling from percentile distributions
  bin_bounds_col = F.expr("transform(slice(VALUE_NUMERIC_ARRAY, 1, size(VALUE_NUMERIC_ARRAY)-1), (x, i) -> struct(x AS LOWER, VALUE_NUMERIC_ARRAY[i+1] AS UPPER))")
  sampled_bin_col = F.element_at(F.shuffle(bin_bounds_col), 1)
  sampled_value_col = F.col("SAMPLED_BIN.LOWER") + F.rand() * (F.col("SAMPLED_BIN.UPPER") - F.col("SAMPLED_BIN.LOWER"))

  generated_df = (
    filtered_meta_df.filter("VALUE_TYPE == 'CONTINUOUS'")
    .join(sample_df, how="inner", on=variable_fields)
    .withColumn("SAMPLED_BIN", sampled_bin_col)  # This needs to be attached using withColumn before sampled_value_col. Trying to do in a single step gives strange results!
    .withColumn("VALUE_NUMERIC", sampled_value_col)
  )

  generated_df = generated_df.select(*sample_df.columns, "VALUE_NUMERIC")

  return generated_df

def create_frequency_field_generator(
  value_type: str,
  *value_fields: Iterable[str]
) -> Callable[[DataFrame, DataFrame, Iterable[str]], DataFrame]:
  """Creates a function that generates a dataframe of (long-formatted) values for fields by sampling from a
  frequency distribution (i.e. of value, weight pairs).

  Args:
      value_type (str): Defines filter on the VALUE_TYPE column in the metadata, which specifies which type of
        metadata this generator operates on (e.g. 'CONTINUOUS', 'DATE', 'DISCRETE')
      *value_fields (Iterable[str]): Names of columns in the metadata which contain the values that the frequency
        distribution is over (each set of values is paired with a weight which defines its frequency)

  Returns:
      Callable[[DataFrame, DataFrame, Iterable[str]], DataFrame]: Generates a dataframe of values by sampling from a
        frequency distribution
  """
  value_type_condition_col = F.col("VALUE_TYPE") == value_type
  value_struct_col = F.struct(*value_fields)

  def _field_generator(meta_df: DataFrame, base_df: DataFrame, variable_fields: Iterable[str]) -> DataFrame:
    """Generates a dataframe of values for fields by sampling from a frequency distribution (i.e. value, weight pairs).

    Args:
        meta_df (pyspark.sql.DataFrame): Metadata for the dataset
        base_df (pyspark.sql.DataFrame): Base data to sample values for, indexed over patients / events.
        variable_fields (Iterable[str]): Names of metadata fields that define the fields to samples (e.g.
          we generally want to sample per combination of 'TABLE_NAME' and 'FIELD_NAME')

    Returns:
        DataFrame: DataFrame of generated data (long-formatted), with a single value per
          unique combination of variable_fields.
    """
    filtered_meta_df = meta_df.filter(value_type_condition_col)
    sample_df = generate_empty_samples(filtered_meta_df, base_df, variable_cols=variable_fields)

    filtered_meta_df = filtered_meta_df.withColumn("VALUE_STRUCT", value_struct_col)

    generated_df = sample_cdf(
      filtered_meta_df,
      sample_df,
      value_field="VALUE_STRUCT",
      frequency_field="WEIGHT",
      variable_fields=variable_fields
    )

    # Select the sampled values, explode the struct after sampling
    generated_df = generated_df.selectExpr(*sample_df.columns, "VALUE_STRUCT.*")

    return generated_df

  return _field_generator

categorical_field_generator = create_frequency_field_generator("CATEGORICAL", "VALUE_STRING")
date_field_generator        = create_frequency_field_generator("DATE", "VALUE_DATE")
discrete_field_generator    = create_frequency_field_generator("DISCRETE", "VALUE_NUMERIC")

def field_generator(
  meta_df: DataFrame,
  base_df: DataFrame,
  pivot_field: str,
  other_variable_fields: Iterable[str]=tuple()
) -> DataFrame:
  """Generates a dataframe of (wide-formatted) values for fields by sampling from metadata.

  Args:
      meta_df (pyspark.sql.DataFrame): Metadata for the dataset
      base_df (pyspark.sql.DataFrame): Base data to sample values for, indexed over patients / events.
      pivot_field (str): Name of the field in the metadata that we want to pivot the resulting generated
        dataframe by. Generally the values in this field are the names of the fields in the dataset
        the artificial data represents
      other_variable_fields (Iterable[str], optional): Fields that, in addition to the pivot field, define
        the unique fields in the metadata. Defaults to tuple().

  Returns:
      DataFrame: DataFrame of generated data (wide-formatted), with a single value per
        unique combination of pivot_field and other_variable_fields.
  """
  filtered_meta_df = meta_df.filter(F.col("VALUE_TYPE").isin(["CATEGORICAL", "DISCRETE", "DATE", "CONTINUOUS"]))

  variable_fields = [pivot_field, *other_variable_fields]
  generator_args = (filtered_meta_df, base_df, variable_fields)
  generated_long_df = union_dfs([
    categorical_field_generator(*generator_args).select(*variable_fields, *base_df.columns, F.col("VALUE_STRING") .alias("VALUE")),
    date_field_generator(       *generator_args).select(*variable_fields, *base_df.columns, F.col("VALUE_DATE")   .alias("VALUE")),
    discrete_field_generator(   *generator_args).select(*variable_fields, *base_df.columns, F.col("VALUE_NUMERIC").alias("VALUE")),
    continuous_field_generator( *generator_args).select(*variable_fields, *base_df.columns, F.col("VALUE_NUMERIC").alias("VALUE")),
  ])

  # Pivot to wide format
  pivot_field_values = get_pivot_field_values(filtered_meta_df, pivot_field)
  generated_wide_df = (
    generated_long_df
    .groupBy(*base_df.columns)
    .pivot(pivot_field, pivot_field_values)
    .agg(F.first("VALUE"))
  )

  return generated_wide_df


# list of available flavours
# cost of each flavour
# quantity - what flavour, how many each ?
# availablility of flavour
# offers
# ice cream combos

# nouns / classes - Flavour, IceCreamShop,
#                   Offers, Menu (Combos),
#                   Bowl(bowl_size), Scoop

class Scoop:
    size: int  # quantity in grams
    scoop_type: str

    def __init__(self, size, scoop_type) -> None:
        self.size = size
        self.scoop_type = scoop_type

class Bowl:
    size: int # number of scoops it can hold
    shape: str # shape of the bowl
    contents = [] # contents of the bowl

    def __init__(self, size, shape) -> None:
        self.size = size
        self.shape = shape

    def add_scoop(self, scoop):
        # add scoop to the content
        self.contents.append(scoop)

class Flavour:
    name: str
    quantity: int  # quantity in grams
    price: int   # price per scoop

    def __init__(self, name, quantity, price):
        self.name = name
        self.quantity = quantity
        self.price = price

    def __repr__(self) -> str:
        return self.name

    def details(self):
        print(f'{self.name} Rs. {self.price} remaining: {self.quantity} grams')

class Menu:
    # price
    # list of ice creams
    ice_creams = []
    combos = []

    def __init__(self) -> None:
        flavour_1 = Flavour('vanilla', 1000, 60)
        flavour_2 = Flavour('chocolate', 1000, 60)
        flavour_3 = Flavour('strawberry', 1000, 60)
        flavour_4 = Flavour('butterscotch', 1000, 60)
        flavour_5 = Flavour('tender-coconut', 1000, 60)
        self.ice_creams = [flavour_1,flavour_2, flavour_3, flavour_4, flavour_5]
        self.combos = [
            [flavour_1, flavour_2, flavour_5],
            [flavour_1, flavour_3],
            [flavour_2, flavour_3, flavour_5],
            [flavour_1, flavour_2, flavour_1]
        ]

    def display(self):
        for ice_cream in self.ice_creams:
            ice_cream.details()

    def display_combos(self):
        for combo in self.combos:
            print(combo)

class IceCreamShop:
    menu = Menu()

    def welcome(self):
        print('Hello, Welcome to our ice cream shop')
        choice = input('Do you want to see the menu? y/n >> ')
        if choice == 'y':
            self.menu.display()
        elif choice == 'n':
            print('maybe next time')
        else:
            print('We dont speak that language')

shop = IceCreamShop()
shop.welcome()


#---------------------------------------------------------------------------------------------------------------------#
# Comfyroll Custom Nodes by RockOfFire and Akatsuzi       https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes
# for ComfyUI                                             https://github.com/comfyanonymous/ComfyUI
#---------------------------------------------------------------------------------------------------------------------#

import random
import string
from .graphics_functions import random_hex_color, random_rgb
from ..categories import icons

#---------------------------------------------------------------------------------------------------------------------#
# Random values
#---------------------------------------------------------------------------------------------------------------------#
class CR_RandomHexColor:

    @classmethod
    def INPUT_TYPES(cls):

        return {"required": {"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),}}

    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", )
    RETURN_NAMES = ("hex_color1", "hex_color2", "hex_color3", "hex_color4", "show_help", )
    FUNCTION = "get_colors"
    CATEGORY = icons.get("Comfyroll/Utils/Random")

    def get_colors(self, seed):

        # Set the seed
        random.seed(seed)

        hex_color1 = random_hex_color()
        hex_color2 = random_hex_color()
        hex_color3 = random_hex_color()
        hex_color4 = random_hex_color()

        show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/Other-Nodes#cr-random-hex-color"

        return (hex_color1, hex_color2, hex_color3, hex_color4, show_help, )

#---------------------------------------------------------------------------------------------------------------------#
class CR_RandomRGB:

    @classmethod
    def INPUT_TYPES(cls):

        return {"required": {"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),}}

    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", )
    RETURN_NAMES = ("rgb_1", "rgb_2", "rgb_3", "rgb_4", "show_help", )
    FUNCTION = "get_colors"
    CATEGORY = icons.get("Comfyroll/Utils/Random")

    def get_colors(self, seed):

        # Set the seed
        random.seed(seed)

        rgb_1 = random_rgb()
        rgb_2 = random_rgb()
        rgb_3 = random_rgb()
        rgb_4 = random_rgb()

        show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/Other-Nodes#cr-random-rgb"

        return (rgb_1, rgb_2, rgb_3, rgb_4, show_help, )

#---------------------------------------------------------------------------------------------------------------------#
class CR_RandomMultilineValues:

    @classmethod
    def INPUT_TYPES(cls):

        types = ["binary", "decimal", "hexadecimal", "alphabetic", "alphanumeric"]

        return {"required": {"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                             "value_type": (types,),
                             "rows": ("INT", {"default": 5, "min": 1, "max": 2048}),
                             "string_length": ("INT", {"default": 5, "min": 1, "max": 2048}),
               }
        }

    RETURN_TYPES = ("STRING", "STRING", )
    RETURN_NAMES = ("multiline_text", "show_help", )
    FUNCTION = "generate"
    CATEGORY = icons.get("Comfyroll/Utils/Random")

    def generate(self, value_type, rows, string_length, seed):

        # Set the seed
        random.seed(seed)

        if value_type == "binary":
            choice_str = '01'
        elif value_type == "decimal":
            choice_str = '0123456789'
        elif value_type == "hexadecimal":
            choice_str = '0123456789abcdef'
        elif value_type == "alphabetic":
            choice_str = string.ascii_letters
        elif value_type == "alphanumeric":
            choice_str = string.ascii_letters + string.digits

        multiline_text = '\n'.join([''.join(random.choice(choice_str) for _ in range(string_length)) for _ in range(rows)])

        show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/Other-Nodes#cr-random-multiline-values"

        return (multiline_text, show_help, )

#---------------------------------------------------------------------------------------------------------------------#

class CR_RandomRGBGradient:

    @classmethod
    def INPUT_TYPES(cls):

        return {"required": {"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                             "rows": ("INT", {"default": 5, "min": 1, "max": 2048}),
               }
        }

    RETURN_TYPES = ("STRING", "STRING", )
    RETURN_NAMES = ("multiline_text", "show_help", )
    FUNCTION = "generate"
    CATEGORY = icons.get("Comfyroll/Utils/Random")

    def generate(self, rows, seed):

        # Set the seed
        random.seed(seed)

        temp = 0
        multiline_text = ""

        for i in range(1, rows + 1):
            print(temp)
            if temp <= 99 - rows + i:
                upper_bound = min(99, temp + (99 - temp) // (rows - i + 1))
                current_value = random.randint(temp, upper_bound)
                multiline_text += f'{current_value}:{random.randint(0, 255)},{random.randint(0, 255)},{random.randint(0, 255)}\n'
                print(multiline_text)
                temp = current_value + 1

        show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/Other-Nodes#cr-random-RGB-gradient"

        return (multiline_text, show_help, )

#---------------------------------------------------------------------------------------------------------------------#
# MAPPINGS
#---------------------------------------------------------------------------------------------------------------------#
# For reference only, actual mappings are in __init__.py
'''
NODE_CLASS_MAPPINGS = {
    # Random
    "CR Random Hex Color": CR_RandomHexColor,
    "CR Random RGB": CR_RandomRGB,
    "CR Random Multiline Values": CR_RandomMultilineValues,
    "CR Random RGB Gradient": CR_RandomRGBGradient,
}
'''


import sys
import numpy as np
import os
import argparse
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
import MDAnalysis as mda
from MDAnalysis.analysis import rms, align, dihedrals, contacts, distances, \
                                hydrogenbonds
from MDAnalysis.analysis.distances import distance_array
import pandas as pd

def main(argv):

    try:
        parser = argparse.ArgumentParser()

        parser.add_argument("-c", "--conform",
                            action = "store",
                            dest = "conform",
                            default = "holo",
                            help = """Chose a conformer for analysis. I.e. "holo" or "apo".""")
        parser.add_argument("-r", "--recalc",
                            action = "store_true",
                            dest = "recalc",
                            default = False,
                            help = """Chose whether the trajectory arrays should  be recomputed.""")
        args = parser.parse_args()

    except argparse.ArgumentError:
        print("Command line arguments are ill-defined, please check the arguments")
        raise

    # Assign group selection from argparse
    conform = args.conform
    recalc = args.recalc

    path_head = "/home/lf1071fu/project_b3"
    fig_path = f"/home/lf1071fu/project_b3/figures/{ conform }"
    str_path = f"{ path_head }/structures"
    if conform == "holo":
        data_path = f"{ path_head }/simulate/holo_conf/data"
    elif conform == "apo":
        data_path = f"{ path_head }/simulate/apo_conf/initial_10us"
    else:
        print("ERROR: chose a valid conform from command line analysis.")
        sys.exit(1)

    #TO DO: CHECK THESE!
    traj_name = "fitted_traj"

    # Load in universe objects for the simulation and the reference structures
    u = mda.Universe(f"{ data_path }/topol.top", f"{ data_path }/{ traj_name }.xtc",
                     topology_format="ITP")

    holo_state = mda.Universe(f"{ str_path }/holo_state.pdb")
    apo_state = mda.Universe(f"{ str_path }/apo_state1.pdb")
    # Starting structure as the ref state
    ref_state = u.select_atoms("protein")

    # Find the core residues with low rmsf for both starting conformations
    if not os.path.exists(f"{ data_path }/core_res.npy"):
        top = f"{ data_path }/topol.top"
        a = mda.Universe(top, f"{ path_head }/simulate/holo_conf/data/full_holo_apo.xtc",
                         topology_format="ITP")
        calphas, rmsf = get_rmsf(a, top, data_path)
        core_res = calphas[(rmsf < 1.5)]
        np.save(f"{ data_path }/core_res.npy", core_res)
    else:
        core_res = np.load(f"{ data_path }/core_res.npy")

    # Atom group selection strings
    beta_flap_group = "backbone and (resid 195-218 or resid 739-762)"
    alpha_flap_group = "backbone and (resid 219-231 or resid 763-775)"
    combo_group = "backbone and (resid 195-231 or  resid 739-775)"
    aln_str = "protein and name CA and ("
    core_holo = [f"resid {i} or " for i in core_res]
    core_apo = [f"resid {i + 544} or " for i in core_res]
    core = aln_str + "".join((core_holo + core_apo))[:-4] + ")"

    # Store calculated outputs as numpy arrays
    np_files = [f"{ data_path }/rmsd_holo.npy", f"{ data_path }/rmsd_apo.npy",
                f"{ data_path }/rmsf.npy", f"{ data_path }/calphas.npy",
                f"{ data_path }/rad_gyration.npy",
                f"{ data_path }/rama.npy", f"{ data_path }/chi_1s.npy",
                f"{ data_path }/salt.npy", f"{ data_path}/salt_dist.npy",
                f"{ data_path }/hpairs.npy"] #, f"{ data_path }/hbonds.npy",
                #f"{ data_path }/hbond_count.npy", f"{ data_path }/times.npy"]

    if all(list(map(lambda x : os.path.exists(x), np_files))) and not recalc:

        R_holo = np.load(np_files[0], allow_pickle=True)
        R_apo = np.load(np_files[1], allow_pickle=True)
        rmsf = np.load(np_files[2], allow_pickle=True)
        calphas = np.load(np_files[3], allow_pickle=True)
        r_gyr = np.load(np_files[4], allow_pickle=True)
        rama = np.load(np_files[5], allow_pickle=True)
        # # chi_1s is a dictionay object
        chi_1s = np.load(np_files[6], allow_pickle=True)
        sc = np.load(np_files[7], allow_pickle=True)
        salt_dist = np.load(np_files[8], allow_pickle=True)
        hpairs = np.load(np_files[9], allow_pickle=True)

        # hbonds = np.load(np_files[7], allow_pickle=True)
        # hbond_count = np.load(np_files[8], allow_pickle=True)
        # times = np.load(np_files[9], allow_pickle=True)

    else:

        u.transfer_to_memory()

        align.AlignTraj(u, ref_state, select=core, in_memory=True).run()

        protein = u.select_atoms("protein")
        with mda.Writer(f"{ data_path }/snap_shots.pdb", protein.n_atoms) as W:
            for ts in u.trajectory:
                if ts.time in [1000000, 3000000, 7000000, 9000000]:
                    W.write(protein)

        # Determine RMSD to ref structures
        R_holo = get_rmsd(u, holo_state, core, ["backbone and resid 8:251",
                          beta_flap_group, alpha_flap_group, combo_group])
        R_apo = get_rmsd(u, apo_state, core, ["backbone and resid 8:251",
                          beta_flap_group, alpha_flap_group, combo_group])

        # Determine RMSF by first finding and aligning to the average structure
        calphas, rmsf = get_rmsf(u, f"{ data_path }/protein.gro", core, data_path)

        # Determine the radius of gyration
        r_gyr = get_rgyr(u)

        # Phi and Psi backbone dihedrals
        rama = dihedrals.Ramachandran(u.residues[195:232]).run()
        rama = rama.results.angles

        # Chi1 dihedrals
        chi_1s = dict()
        for res in u.residues[195:232]:
            group = res.chi1_selection()
            if group is not None:
                 dihs = dihedrals.Dihedral([group]).run()
                 chi_1s[res.ix] = dihs.results.angles

        # Salt-bridge contacts
        sc = get_salt_contacts(u, ref_state, data_path)
        salt_dist = get_bridges(u, data_path)

        # # Hydrogen bonds analysis (intramolecular, not with solvent)
        # hbonds, hbond_count, times = get_hbonds(u, path)
        hpairs = get_hbond_pairs(u, data_path)

        for f, v in zip(np_files, [R_holo, R_apo, rmsf, calphas, r_gyr, rama, chi_1s,
                                   sc, salt_dist, hpairs]): #hbonds, hbond_count, times]):
            np.save(f, v)

    # Make plots for simulation analysis
    # plot_rmsd(R_holo, R_apo, fig_path)
    # plot_rmsf(calphas, rmsf, fig_path)
    # plot_rmsd_time(R_apo, fig_path)
    # plot_rgyr(r_gyr, fig_path)

    # Plots for individual residues: Ramachandran and Chi_1
    # for i, id in enumerate(np.arange(198,232)):
    #     res = u.residues[id]
    #     plot_ramachandran(rama[:,i,0], rama[:,i,1], res, fig_path)
    # for key, chi_1 in chi_1s.items():
    #     res = u.residues[key]
    #     plot_chi1(res, chi_1, fig_path)
    #
    print(hpairs.shape)
    # plot_salt_bridges(sc, fig_path)
    # plot_bridges(salt_dist, r_gyr, fig_path)
    plot_hbond_pairs(hpairs, r_gyr, fig_path)
    #plot_hbonds_count(hbond_count, times, fig_path)

    return None

def get_rmsd(system, reference, alignment, group):
    """Determines the rmsd over the trajectory against a reference structure.

    The MDAnalysis.analysis.rms.results array is saved as a numpy array file,
    which can be loaded if it has alreay been determined.

    Parameters
    ----------
    system : MDAnalysis.core.universe
        The universe object, with the trajectory loaded in.
    reference : MDAnalysis.core.universe
        The topology of the reference structure.
    alignment : MDAnalysis.core.groups.AtomGroup
        The group for the alignment, i.e. all the alpha carbons.
    group : MDAnalysis.core.groups.AtomGroup
        The group for the RMSD, e.g. the beta flap residues.

    Returns
    -------
    rmsd_arr : np.ndarray
        A timeseries of the RMSD against the given reference, for the given
        atom group.
    """
    if type(group) != list:
        group = [group]
    R = rms.RMSD(system,
                 reference,  # reference universe or atomgroup
                 select=alignment,  # group to superimpose and calculate RMSD
                 groupselections=group)  # groups for RMSD
    R.run()

    rmsd_arr = R.results.rmsd

    return rmsd_arr

def get_rmsf(u, top, ref_group, path):
    """Determines the RMSF per residue.

    The average structure is calculated and then used as the reference for
    structure alignment, before calculating the rmsf per residue.

    Parameters
    ----------
    u : MDAnalysis.core.universe
        The universe object, with the trajectory loaded in.
    top : str
        The topology file, e.g. xxx.gro.
    ref_group : str
        Atom group selection string of the core residues for the alignment.
    path : str
        Path to the primary working directory.

    Returns
    -------
    resnums : MDAnalysis.core.groups.AtomGroup
        The atom group used for the RMSF calculation consists of all the alpha
        carbons in the protein.
    rmsf : np.ndarray
        The rmsf of the selected atom groups.

    """
    # Perform alignment to the average structure and write to a new traj file.
    protein = u.select_atoms("protein")
    prealigner = align.AlignTraj(u, u, select="protein and name CA",
                                 in_memory=True).run()
    ref_coords = u.trajectory.timeseries(asel=protein).mean(axis=1)
    ref = mda.Merge(protein).load_new(ref_coords[:, None, :],
                              order="afc")
    aligner = align.AlignTraj(u, ref, select=ref_group, in_memory=True).run()
    with mda.Writer(f"{ path }/rmsfit.xtc", n_atoms=u.atoms.n_atoms) as W:
        for ts in u.trajectory:
            W.write(u.atoms)

    from MDAnalysis.analysis.rms import RMSF

    u = mda.Universe(top, f"{ path }/rmsfit.xtc")
    calphas = protein.select_atoms("protein and name CA")

    rmsfer = RMSF(calphas).run()

    return calphas.resnums, rmsfer.results.rmsf

def get_salt_contacts(u, ref_state, path):
    """Runs the Contacts module to get the number of salt contacts.

    Arginine and lysine residues are considered as the basic residues, while
    aspargine and glutamine are classified as acidic residues. A soft cut off
    is used for determining contacts within a radius of 4.5 A.

    Parameters
    ----------
    u : MDAnalysis.core.universe.Universe
        The universe object, with the trajectory loaded in.
    ref_state : MDAnalysis.core.universe.Universe
        The universe object for the reference state.
    path : str
        Path to the primary working directory.

    Returns
    -------
    sc : np.ndarray
        A timeseries of the fraction of retained salt bridges, with respect to
        the reference state.
    """
    sel_basic = "(resname ARG LYS) and (name NH* NZ*)"
    sel_acidic = "(resname ASP GLU) and (name OE* OD*)"
    acidic = ref_state.select_atoms(sel_acidic)
    basic = ref_state.select_atoms(sel_basic)
    salt_contacts = contacts.Contacts(u, select=(sel_acidic, sel_basic),
                    refgroup=(acidic, basic), radius=4.5, method="soft_cut")

    salt_contacts.run()
    sc = salt_contacts.results.timeseries

    return sc

def get_bridges(u, path):
    """Calculates the distance of special salt contacts in the alpha + beta flaps.

    Arginine and lysine residues are considered as the basic residues, while
    aspargine and glutamine are classified as acidic residues.

    Parameters
    ----------
    u : MDAnalysis.core.universe.Universe
        The universe object, with the trajectory loaded in.
    path : str
        Path to the primary working directory.

    Returns
    -------
    distances : np.ndarray
        Trajectory of the four distance arrays.

    """
    acid_base_pairs = [
        (("resid 208 and name NH*"),("resid 222 and name OE*")), # R208 -- E222
        (("resid 227 and name NZ*"),("resid 213 and name OD*")), # K227 -- D213
        #(("resid 232 and name NZ*"),("resid 149 and name OE*")), # K232 -- E149
        (("resid 202 and name NH*"),("resid 210 and name OE*")), # R202 -- E210
        (("resid 57 and name NZ*"),("resid 200 and name OE*"))] # K57 -- E200

    pairs = []

    for b, a in acid_base_pairs:
        sel_basic = u.select_atoms(b)
        sel_acidic = u.select_atoms(a)
        pairs.append((sel_basic, sel_acidic))

        #dist_pair1 = distance_array(u.coord[ca_pair1_1], u.coord[ca_pair1_2])
    distances = np.zeros((u.trajectory.n_frames, 4))

    # Loop over all frames in the trajectory
    for ts in u.trajectory:
        # Calculate the distances between the four acid-base pairs for this frame
        d1 = distance_array(pairs[0][0].positions, pairs[0][1].positions)
        d2 = distance_array(pairs[1][0].positions, pairs[1][1].positions)
        d3 = distance_array(pairs[2][0].positions, pairs[2][1].positions)
        d4 = distance_array(pairs[3][0].positions, pairs[3][1].positions)
        # Store the distances in the distances array
        distances[ts.frame] = [np.min(d1), np.min(d2), np.min(d3), np.min(d4)]

    return distances

def get_hbonds(u, path):
    """Find the intramolecular hydrogen bonds in the simulation.

    The donor-acceptor distance cutoff is 3 A and the donor-hydrogen-acceptor
    angle cutoff is minimum 150 deg. The function generates two files, one to
    track the types of hbonds formed and one to identify significant and
    persistent bonds.

    Parameters
    ----------
    u : MDAnalysis.core.universe
        The universe object, with the trajectory loaded in.
    path : str
        Path to the primary working directory.

    Returns
    -------
    hbonds : MDAnalysis...HydrogenBondAnalysis
        Each row consists of a particular h-bond observation: [frame,
        donor_index, hydrogen_index, acceptor_index, DA_distance, DHA_angle].
    hbond_count : np.ndarray
        The hbond count as a trajectory.

    """
    hbonds = hydrogenbonds.hbond_analysis.HydrogenBondAnalysis(universe=u)
    hbonds.hydrogens_sel = hbonds.guess_hydrogens("protein")
    hbonds.acceptors_sel = hbonds.guess_acceptors("protein")
    hbonds.run()
    tau_timeseries, timeseries = hbonds.lifetime()

    # these get their own np arrays
    times = hbonds.times
    hbond_count = hbonds.count_by_time()

    # Analyze the bond types
    f = open(f"{ path }/hbond_types.txt", "w")
    f.write("The average number of each type of hydrogen bond formed at each"
            " frame:\n")

    for donor, acceptor, count in hbonds.count_by_type():

        donor_resname, donor_type = donor.split(":")
        n_donors = u.select_atoms(f"resname { donor_resname } " \
                                  f"and type { donor_type }").n_atoms

        # average number of hbonds per donor molecule per frame
        # multiply by two as each hydrogen bond involves two water molecules
        mean_count = 2 * int(count) / (hbonds.n_frames * n_donors)
        if mean_count > 0.1:
            f.write(f"{donor} to {acceptor}: {mean_count:.2f}\n")
    f.close()

    f = open(f"{ path }/significant_hbonds.txt", "w")
    f.write("A descending list of h-bonds formed between atom pairs.\n"\
            "Donor, Hydrogen, Acceptor, Count\n")

    def describe_at(atom):
        a = u.atoms[atom-1]
        return f"{ a.resname }{ a.resid }--{ a.type }"

    # List the most important h-bonds by order of occurence
    c = hbonds.count_by_ids()
    for d, h, a, count in c:
        if count < 10:
            continue
        f.write(f"{ describe_at(d) }\t{ describe_at(h) }\t{ describe_at(a) } "\
                f"\t{ count }\n")
    f.close()

    return hbonds, hbond_count, times

def get_hbond_pairs(u, path):
    """Calculates the distance of special hbond contacts in the alpha + beta flaps.

    Parameters
    ----------
    u : MDAnalysis.core.universe.Universe
        The universe object, with the trajectory loaded in.
    path : str
        Path to the primary working directory.

    Returns
    -------
    distances : np.ndarray
        Trajectory of the four distance arrays.

    """
    hbond_pairs = [
        (("resid 204 and name OD1"),("(resid 206 or resid 208) and name H")),
        (("resid 231 and name O*"),("resid 209 and name NH*")),
        (("resid 197 and name OD1"),("resid 209 and name N*")),
        (("resid 53 and name H*"),("resid 200 and name O*"))]

    pairs = []

    for a, d in hbond_pairs:
        sel_accept = u.select_atoms(a)
        sel_donor = u.select_atoms(d)
        pairs.append((sel_accept, sel_donor))

    distances = np.zeros((u.trajectory.n_frames, len(pairs)))

    # Loop over all frames in the trajectory
    for ts in u.trajectory:
        # Calculate the distances between the 3 hbond pairs for this frame
        d1 = distance_array(pairs[0][0].positions, pairs[0][1].positions)
        d2 = distance_array(pairs[1][0].positions, pairs[1][1].positions)
        d3 = distance_array(pairs[2][0].positions, pairs[2][1].positions)
        d4 = distance_array(pairs[3][0].positions, pairs[3][1].positions)

        # Store the distances in the distances array
        distances[ts.frame] = [np.min(d1), np.min(d2), np.min(d3), np.min(d3)]

    return distances

def get_rgyr(u):
    """
    """
    r_gyr = []
    protein = u.select_atoms("protein")
    for ts in u.trajectory:
       r_gyr.append((u.trajectory.time, protein.radius_of_gyration()))
    r_gyr = np.array(r_gyr)
    return r_gyr

def plot_rmsd(r_holo, r_apo, path):
    """Makes a plot of the rmsd against two ref structures with a color bar.

    Parameters
    ----------
    r_holo : MDAnalysis.analysis.rms.RMSD or nparray
        A time series of the rmsd with the holo crystal structure as reference.
    r_apo : MDAnalysis.analysis.rms.RMSD or nparray
        A time series of the rmsd with the apo crystal structure as reference.
    path : str
        Path to the primary working directory.

    Returns
    -------
    None.

    """
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,8))

    d = ax.scatter(r_holo[:,3], r_apo[:,3], c=np.arange(0,len(r_holo[:,0])),
                cmap="cividis", label=r"Full backbone", marker="X")
    d = ax.scatter(r_holo[:,4], r_apo[:,4], c=np.arange(0,len(r_holo[:,0])),
                cmap="cividis", label=r"$\beta$-flap", marker="o")
    d = ax.scatter(r_holo[:,5], r_apo[:,5], c=np.arange(0,len(r_holo[:,0])),
                cmap="cividis", label=r"$\alpha$-flap", marker="D")

    # Colormap settings
    cbar = plt.colorbar(d)
    cbar.set_label(r'Time [$\mu$s]', fontsize=28, labelpad=10)
    cbar.ax.yaxis.set_ticks(np.arange(0,len(r_holo[:,0]),1000))
    cticks = list(map(lambda x: str(x/1000).split(".")[0],
                      np.arange(0,len(r_holo[:,0]),1000)))
    cbar.ax.yaxis.set_ticklabels(cticks)
    cbar.ax.tick_params(labelsize=16, direction='out', width=2, length=5)
    cbar.outline.set_linewidth(2)

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"RMSD to holo state ($\AA$)", labelpad=5, \
                    fontsize=28)
    ax.set_ylabel(r"RMSD to apo state ($\AA$)", labelpad=5, fontsize=28)
    _, xmax = ax.get_xlim()
    _, ymax = ax.get_ylim()
    ax.set_xlim(0,xmax)
    ax.set_ylim(0,ymax)

    plt.legend(fontsize=24)

    plt.savefig(f"{ path }/rmsd.pdf")
    plt.close()

    return None

def plot_rmsf(calphas, rmsf, path):
    """Makes an RMSF plot.

    Parameters
    ----------
    resnums : MDAnalysis.core.groups.AtomGroup
        The atom group used for the RMSF calculation consists of all the alpha
        carbons in the protein.
    rmsf : np.ndarray
        The rmsf of the selected atom groups.
    path : str
        Path to the primary working directory.

    Returns
    -------
    None.

    """
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,4))
    resids = list(map(lambda x : x + 544, calphas))
    plt.plot(calphas, rmsf, lw=3, color="#47abc4")

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"Residue number", labelpad=5, fontsize=28)
    ax.set_ylabel(r"RMSF ($\AA$)", labelpad=5, fontsize=28)
    bottom, top = ax.get_ylim()
    ax.vlines([195,217.5], bottom, top, linestyles="dashed", lw=3,
              colors="#FF1990")
    ax.vlines([219.5,231], bottom, top, linestyles="dashed", lw=3,
              colors="#dba61f")
    ax.set_ylim(-1,10)

    plt.savefig(f"{ path }/rmsf.pdf")
    plt.close()

    return None

def plot_rmsd_time(r_apo, path):
    """
    """
    fig, ax = plt.subplots(constrained_layout=True, figsize=(24,8))
    time = r_apo[:,0]
    #plt.plot(time, r_apo[:,2], lw=2, color="#02ab64", alpha=0.8,
    #         label="core")
    plt.plot(time[::10], r_apo[:,5][::10], lw=3, color="#02ab64", alpha=0.8,
             label="alpha flap")
    plt.plot(time[::10], r_apo[:,4][::10], lw=3, color="#dba61f", alpha=0.8,
             label="beta flap")
    plt.plot(time[::10], r_apo[:,3][::10], lw=3, color="#47abc4", alpha=0.8,
             label="full backbone")

    # Plot settings
    ax.tick_params(axis='y', labelsize=32, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=32, direction='in', width=2, \
                    length=5, pad=10)
    xticks = np.arange(0,len(time),1000)
    ax.set_xticks(xticks)
    ax.set_xticklabels(list(map(lambda x : str(x/1000).split(".")[0], xticks)))
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"time ($\mu s$)", labelpad=5, fontsize=38)
    ax.set_ylabel(r"RMSD ($\AA$)", labelpad=5, fontsize=38)
    _, ymax = ax.get_ylim()
    ax.set_ylim(0,ymax)
    plt.legend(fontsize=28)

    plt.savefig(f"{ path }/rmsd_time.pdf", dpi=300)
    plt.close()

    return None

def plot_ramachandran(phis, psis, res, path):
    """Make a Ramachandran plot for the given residues.

    Parameters
    ----------
    dihs_phi : np.ndarray
        A timeseries of the phi dihedrals for the residue group.
    dihs_psi : np.ndarray
        A timeseries of the psi dihedrals for the residue group.
    resnum : MDAnalysis.core.groups.Residue
        The residue object for the ramachandran plot.
    path : str
        Path to the primary working directory.

    Returns
    -------
    None.

    """
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,8))

    d = ax.scatter(phis, psis, c=np.arange(0,len(phis)), cmap="cividis",
                   label=f"{ res.resname } { res.resid }")

    # Colormap settings
    cbar = plt.colorbar(d)
    cbar.set_label(r'Time [$\mu$s]', fontsize=28, labelpad=10)
    cbar.ax.yaxis.set_ticks(np.arange(0,len(psis),1000))
    cticks = list(map(lambda x: str(x/1000).split(".")[0],
                      np.arange(0,len(psis),1000)))
    cbar.ax.yaxis.set_ticklabels(cticks)
    cbar.ax.tick_params(labelsize=16, direction='out', width=2, length=5)
    cbar.outline.set_linewidth(2)

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"$\Phi$", labelpad=5, \
                    fontsize=28)
    ax.set_ylabel(r"$\Psi$", labelpad=5, fontsize=28)
    plt.xlim([-180, 180])
    plt.ylim([-180, 180])
    plt.legend(fontsize=20)

    # Save plot to file
    plot_path = f"{ path }/ramachandran"
    if not os.path.exists(plot_path):
       os.makedirs(plot_path)
    plt.savefig(f"{ plot_path }/res_{ res.resid }.pdf")
    plt.close()

    return None

def plot_chi1(res, chi_1, path):
    """Make a chi_1 plot for the given residue.

    Parameters
    ----------
    res : MDAnalysis.core.groups.Residue
        The residue object for the chi_1 plot.
    chi_1 : np.ndarray
        A timeseries of the chi 1 dihedrals for the residue group.
    path : str
        Path to the primary working directory.

    Returns
    -------
    None.

    """
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,8))

    d = ax.hist(chi_1, bins=180, density=True, color="#b21856",
                label=f"{ res.resname } { res.resid }")

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"$\chi_1$ (°)", labelpad=5, \
                    fontsize=28)
    ax.set_ylabel(r"Density", labelpad=5, fontsize=28)
    plt.xlim([-180, 180])
    plt.legend(fontsize=20)

    plot_path = f"{ path }/chi_1s"
    if not os.path.exists(plot_path):
       os.makedirs(plot_path)
    plt.savefig(f"{ plot_path }/res_{ res.resid }.pdf")
    plt.close()

    return None

def plot_salt_bridges(sc, path):
    """Make a plot for the fraction of salt contacts compared to the reference.

    A soft cut-off from 4 A was used to determine the fraction of salt bridges
    retained from the reference structure. This includes bridges between the
    charged residues, i.e. N in ARG and LYS with O in ASP and GLU.

    Parameters
    ----------
    sc : np.ndarray
        A timeseries of the fraction of salt bridge contacts, relative to the
        reference structure.
    path : str
        Path to the primary working directory.

    Returns
    -------
    None.

    """

    print(sc.shape)
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,8))
    ave_contacts = np.mean(sc[:, 1])
    print(f"average contacts = { ave_contacts }")
    sc = sc[::40]

    ax.scatter(sc[:, 0], sc[:, 1], s=150, color="#00C27B")

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"Frame", labelpad=5, fontsize=24)
    ax.set_ylabel(r"Fraction of Native Contacts", labelpad=5, fontsize=24)
    plt.ylim([0, 1])

    plt.savefig(f"{ path }/salt_contacts.pdf")
    plt.close()

    return None

def plot_hbonds_count(hbond_count, times, path):
    """Make a plot for the number of hbonds at each time step.

    Parameters
    ----------
    hbond_count : np.ndarray
        A trajectory of the hydrogen bond count.
    times : np.ndarray
        The trajectory times as an array in ps.
    path : str
        Path to the primary working directory.

    Returns
    -------
    None.

    """
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,8))

    plt.scatter(times[::100]/(1e6), hbond_count[::100], s=150,
                color="#00C27B")

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"Time ($\mu$s)", labelpad=5, fontsize=24)
    ax.set_ylabel(r"$N_{HB}$", labelpad=5, fontsize=24)
    plt.ylim([0,100])

    plt.savefig(f"{ path }/hbonds_count.pdf")
    plt.close()

    return None

def plot_rgyr(r_gyr, path):
    fig, ax = plt.subplots(constrained_layout=True, figsize=(24,8))

    time = r_gyr[:,0] / 1000
    plt.plot(time[::10], r_gyr[:,1][::10], "--", lw=3, color="#02ab64", label=r"$R_G$")

    # Plot settings
    ax.tick_params(axis='y', labelsize=32, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=32, direction='in', width=2, \
                    length=5, pad=10)
    xticks = np.arange(0,len(time),1000)

    ax.set_xticks(xticks)
    x_labels = list(map(lambda x : str(x/1000).split(".")[0], xticks))
    ax.set_xticklabels(x_labels)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_xlabel(r"time ($\mu s$)", labelpad=5, fontsize=38)
    ax.set_ylabel(r"Radius of Gyration $R_G$ ($\AA$)", labelpad=5, fontsize=38)
    plt.legend(fontsize=28)

    plt.savefig(f"{ path }/rad_gyration.pdf", dpi=300)
    plt.close()

    return None

def plot_bridges(salt_dist, r_gyr, path):
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,4))

    time = r_gyr[:,0] / 1000
    stride = 10
    colors = ["#1f77b4", "#ff7f0e", "#2ca02c", "#9467bd"]
    labels = [r"$\alpha - \beta \:$ R208$-$E222",
              r"$\alpha - \beta \:$ K227$-$D213",
              #r"$\alpha - pocket \:$ K232-E149",
              r"$\beta - \beta \:$ R202$-$E210",
              r"Pocket$-\beta$ K57$-$E200"]

    for i in range(4):

        plt.plot(time[::stride], salt_dist[::stride,i], "-", lw=3, color=colors[i],
            label=labels[i], alpha=0.8,
            path_effects=[pe.Stroke(linewidth=5, foreground='#595959'), pe.Normal()])

    ax.axhline(y=4.5, linestyle='--', lw=3, color='red', alpha=0.5,
        path_effects=[pe.Stroke(linewidth=5, foreground='#595959'), pe.Normal()])

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    xticks = np.arange(0,len(time),1000)
    ax.set_xticks(xticks)
    ax.set_xticklabels(list(map(lambda x : str(x/1000).split(".")[0], xticks)))
    ax.set_xlabel(r"Time ($\mu s$)", labelpad=5, fontsize=24)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_ylabel("Distance ($\AA$)", labelpad=5, fontsize=24)
    plt.ylim([0, 40])
    _, xmax = ax.get_xlim()
    ax.set_xlim(0,xmax)
    plt.legend(fontsize=24, ncol=2)

    plt.savefig(f"{ path }/salt_bridges.pdf", dpi=300)
    plt.close()

    return None

def plot_hbond_pairs(hpairs, r_gyr, path):
    fig, ax = plt.subplots(constrained_layout=True, figsize=(12,4))

    time = r_gyr[:,0] / 1000
    stride = 10
    colors = ["#1f77b4", "#ff7f0e", "#2ca02c", "#9467bd"]
    labels = [r"N204$-$R208", r"S231$-$R209",
              r"N197$-$R209", r"N53$-$E200"]

    print(hpairs[::1000,3])

    for i in range(4):

        plt.plot(time[::stride], hpairs[::stride,i], "-", lw=3, color=colors[i],
            label=labels[i], alpha=0.8,
            path_effects=[pe.Stroke(linewidth=5, foreground='#595959'), pe.Normal()])

    ax.axhline(y=4.5, linestyle='--', lw=3, color='red', alpha=0.5,
        path_effects=[pe.Stroke(linewidth=5, foreground='#595959'), pe.Normal()])

    # Plot settings
    ax.tick_params(axis='y', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    ax.tick_params(axis='x', labelsize=18, direction='in', width=2, \
                    length=5, pad=10)
    xticks = np.arange(0,len(time),1000)
    ax.set_xticks(xticks)
    ax.set_xticklabels(list(map(lambda x : str(x/1000).split(".")[0], xticks)))
    ax.set_xlabel(r"Time ($\mu s$)", labelpad=5, fontsize=24)
    for i in ["top","bottom","left","right"]:
        ax.spines[i].set_linewidth(2)
    ax.grid(True)
    ax.set_ylabel("Distance ($\AA$)", labelpad=5, fontsize=24)
    plt.ylim([0, 30])
    _, xmax = ax.get_xlim()
    ax.set_xlim(0,xmax)
    plt.legend(fontsize=24, ncol=2)

    plt.savefig(f"{ path }/hpairs.pdf", dpi=300)
    plt.close()

    return None

if __name__ == '__main__':
    main(sys.argv)


# Code based on an original given in
# https://github.com/agaidus/census_data_extraction/blob/master/census_mapper.py

import geopandas as gpd
import pandas as pd
from shapely.geometry import Point
from numpy.random import RandomState, uniform

def generate_random_points_in_polygon(poly, num_points, seed=None):
	"""
	Create a list of randomly generated points within a polygon.

	Parameters
	----------
	poly : Polygon
	num_points : int
		The number of random points to create within the polygon
	seed : int, optional
		A random seed

	Returns
	-------
	List
	"""
	min_x, min_y, max_x, max_y = poly.bounds
	points = []
	i = 0
	while len(points) < num_points:
		s = RandomState(seed + i) if seed else RandomState(seed)
		random_point = Point([s.uniform(min_x, max_x), s.uniform(min_y, max_y)])
		if random_point.within(poly):
			points.append(random_point)
		i += 1
	return points

def generate_points_in_areas(gdf, values, points_per_unit=1, seed=None):
	"""
	Create a GeoSeries of random points in polygons.

	Parameters
	----------
	gdf : GeoDataFrame
		The areas in which to create points
	values : str or Series
		The [possibly scaled] number of points to create in each area
	points_per_unit : numeric, optional
		The rate to scale the values in point generation.
	seed : int, optional
		A random seed

	Returns
	-------
	GeoSeries
	"""
	geometry = gdf.geometry
	if isinstance(values, str) and values in gdf.columns:
		values = gdf[values]
	new_values = (values / points_per_unit).astype(int)
	g = gpd.GeoDataFrame(data={'vals': new_values}, geometry=geometry)
	a = g.apply(lambda row: tuple(generate_random_points_in_polygon(row['geometry'], row['vals'], seed)), 1)
	s = lambda x: gpd.GeoSeries(x)
	b = gpd.GeoSeries(a.apply(s).stack(), crs=geometry.crs)
	b.name = 'geometry'
	return b


import pyspark
from pyspark.sql import SparkSession
from pyspark.sql.functions import lit
from packages.etl.transformer import Transformer

class YellowTransformer(Transformer):

    def __init__(self, spark: SparkSession):
        super().__init__(spark=spark, entity_name='yellow')

    def apply_transformation(self, df: pyspark.sql.DataFrame):

        # Renomear colunas do DataFrame green
        df = df.withColumnRenamed('tpep_pickup_datetime', 'pickup_datetime') \
               .withColumnRenamed('tpep_dropoff_datetime', 'dropoff_datetime')

        df = df.withColumn('ehail_fee', lit(0))
        df = df.withColumn('trip_type', lit(0))

        df = df.withColumn('source', lit('yellow'))

        df = df.where("(trip_distance <= 30 and trip_distance > 0) "
                      "AND (passenger_count > 0 AND passenger_count < 5) "
                      "AND (pickup_datetime < dropoff_datetime)"
                      "AND (pickup_datetime >= '2018-01-01' and pickup_datetime < '2023-01-01')"
                      "AND (dropoff_datetime >= '2018-01-01' and dropoff_datetime < '2023-01-01')")
        return df

    def execute(self):
        df = self._spark_session.read.parquet(f"{self.filesystem_path}/raw/yellow")
        df_transformed = self.apply_transformation(df=df)
        self._loader.write_to_filesystem(layer='trusted', df=df_transformed)