import tempfile
import unittest
from collections import OrderedDict
from itertools import product

import numpy as np
import parameterized

# isort: off
import torch
import tensorrt as trt
# isort: on
from parameterized import parameterized
from transformers import BertConfig, BertForQuestionAnswering, BertModel

import tensorrt_llm
import tensorrt_llm.runtime
from tensorrt_llm import Builder
from tensorrt_llm._utils import trt_dtype_to_torch
from tensorrt_llm.network import net_guard
from tensorrt_llm.plugin.plugin import ContextFMHAType
from tensorrt_llm.runtime import TensorInfo

def extract_layer_idx(name):
    ss = name.split('.')
    for s in ss:
        if s.isdigit():
            return s
    return None

def split(v, tp_size, idx, dim=0):
    if tp_size == 1:
        return v
    if len(v.shape) == 1:
        return np.ascontiguousarray(np.split(v, tp_size)[idx])
    elif len(v.shape) == 2:
        return np.ascontiguousarray(np.split(v, tp_size, axis=dim)[idx])
    return None

def load_from_hf_bert(tensorrt_llm_bert,
                      hf_bert,
                      hf_bert_config,
                      rank=0,
                      tensor_parallel=1,
                      fp16=False):
    qkv_weight = [[None, None, None]
                  for _ in range(hf_bert_config.num_hidden_layers)]

    qkv_bias = [[None, None, None]
                for _ in range(hf_bert_config.num_hidden_layers)]

    torch_dtype = torch.float16 if fp16 else torch.float32
    for k, v in hf_bert.state_dict().items():
        v = v.to(torch_dtype).cpu().numpy()
        if 'embeddings.word_embeddings.weight' in k:
            tensorrt_llm_bert.embedding.vocab_embedding.weight.value = v
        elif 'embeddings.position_embeddings.weight' in k:
            tensorrt_llm_bert.embedding.position_embedding.weight.value = v
        elif 'embeddings.token_type_embeddings.weight' in k:
            tensorrt_llm_bert.embedding.token_embedding.weight.value = v
        elif 'embeddings.LayerNorm.weight' in k:
            tensorrt_llm_bert.embedding.embedding_ln.weight.value = v
        elif 'embeddings.LayerNorm.bias' in k:
            tensorrt_llm_bert.embedding.embedding_ln.bias.value = v
        else:
            layer_idx = extract_layer_idx(k)
            if layer_idx is None:
                continue
            idx = int(layer_idx)
            if 'attention.output.dense.weight' in k:
                tensorrt_llm_bert.layers[
                    idx].attention.dense.weight.value = split(v,
                                                              tensor_parallel,
                                                              rank,
                                                              dim=1)
            elif 'attention.output.dense.bias' in k:
                tensorrt_llm_bert.layers[idx].attention.dense.bias.value = v
            elif 'attention.output.LayerNorm.weight' in k:
                tensorrt_llm_bert.layers[idx].input_layernorm.weight.value = v
            elif 'attention.output.LayerNorm.bias' in k:
                tensorrt_llm_bert.layers[idx].input_layernorm.bias.value = v
            elif 'intermediate.dense.weight' in k:
                tensorrt_llm_bert.layers[idx].mlp.fc.weight.value = split(
                    v, tensor_parallel, rank)
            elif 'intermediate.dense.bias' in k:
                tensorrt_llm_bert.layers[idx].mlp.fc.bias.value = split(
                    v, tensor_parallel, rank)
            elif 'output.dense.weight' in k:
                tensorrt_llm_bert.layers[idx].mlp.proj.weight.value = split(
                    v, tensor_parallel, rank, dim=1)
            elif 'output.dense.bias' in k:
                tensorrt_llm_bert.layers[idx].mlp.proj.bias.value = v
            elif 'output.LayerNorm.weight' in k:
                tensorrt_llm_bert.layers[idx].post_layernorm.weight.value = v
            elif 'output.LayerNorm.bias' in k:
                tensorrt_llm_bert.layers[idx].post_layernorm.bias.value = v
            elif 'attention.self.query.weight' in k:
                qkv_weight[idx][0] = v
            elif 'attention.self.query.bias' in k:
                qkv_bias[idx][0] = v
            elif 'attention.self.key.weight' in k:
                qkv_weight[idx][1] = v
            elif 'attention.self.key.bias' in k:
                qkv_bias[idx][1] = v
            elif 'attention.self.value.weight' in k:
                qkv_weight[idx][2] = v
            elif 'attention.self.value.bias' in k:
                qkv_bias[idx][2] = v

    for i in range(hf_bert_config.num_hidden_layers):
        tensorrt_llm_bert.layers[i].attention.qkv.weight.value = split(
            np.concatenate(qkv_weight[i]), tensor_parallel, rank)
        tensorrt_llm_bert.layers[i].attention.qkv.bias.value = split(
            np.concatenate(qkv_bias[i]), tensor_parallel, rank)

def load_from_hf_qa_bert(tensorrt_llm_qa_bert,
                         hf_qa_bert,
                         hf_bert_config,
                         rank=0,
                         tensor_parallel=1,
                         fp16=False):
    load_from_hf_bert(tensorrt_llm_qa_bert.bert, hf_qa_bert, hf_bert_config,
                      rank, tensor_parallel, fp16)
    states = hf_qa_bert.state_dict()

    torch_dtype = torch.float16 if fp16 else torch.float32

    tensorrt_llm_qa_bert.qa_outputs.weight.value = states[
        'qa_outputs.weight'].to(torch_dtype).cpu().numpy()
    tensorrt_llm_qa_bert.qa_outputs.bias.value = states['qa_outputs.bias'].to(
        torch_dtype).cpu().numpy()

class TestBert(unittest.TestCase):

    def load_test_cases():
        models = [BertModel.__name__, BertForQuestionAnswering.__name__]
        test_cases = []
        test_cases += product(models, [False], [False], [False],
                              [ContextFMHAType.disabled], ['float32'])
        test_cases += product(models, [False], [True], [True], [
            ContextFMHAType.disabled, ContextFMHAType.enabled,
            ContextFMHAType.enabled_with_fp32_acc
        ], ['float16'])

        return test_cases

    def custom_name_func(testcase_func, param_num, param):
        return "%s_%s" % (
            testcase_func.__name__,
            parameterized.to_safe_name("_".join(str(x) for x in param.args)),
        )

    @parameterized.expand(load_test_cases, name_func=custom_name_func)
    def test_bert(self, model, use_refit, use_plugin, fast_building,
                  context_fmha_type, dtype):
        tensorrt_llm.logger.set_level('error')
        fp16 = (dtype == 'float16')
        world_size = 1
        rank = 0
        batch_size = 8
        input_len = 128
        vocab_size = 51200
        num_layers = 12
        num_heads = 12
        hidden_act = 'gelu'
        max_position_embeddings = 512
        hidden_size = 768
        bs_range = [1, (batch_size + 1) // 2, batch_size]
        inlen_range = [1, (input_len + 1) // 2, input_len]
        torch_dtype = torch.float16 if fp16 else torch.float32
        trt_dtype = trt.float16 if fp16 else trt.float32
        timing_cache = 'model.cache'

        torch.manual_seed(0)

        builder = Builder()
        with tempfile.TemporaryDirectory() as tmpdirname:
            builder_config = builder.create_builder_config(
                name=model,
                precision='float16' if fp16 else 'float32',
                timing_cache=timing_cache,
                tensor_parallel=world_size,  # TP only
                use_refit=use_refit)
            network = builder.create_network()
            if use_plugin:
                network.plugin_config.set_bert_attention_plugin(dtype)
            if fast_building:
                network.plugin_config.set_gemm_plugin(dtype)
            network.plugin_config.set_context_fmha(context_fmha_type)
            with net_guard(network):
                # Prepare inputs
                # TODO: could class be better than dict for profiles?
                input_ids = tensorrt_llm.Tensor(name='input_ids',
                                                dtype=trt.int32,
                                                shape=[-1, -1],
                                                dim_range=OrderedDict([
                                                    ('batch_size', [bs_range]),
                                                    ('input_len', [inlen_range])
                                                ]))
                input_lengths = tensorrt_llm.Tensor(name='input_lengths',
                                                    dtype=trt.int32,
                                                    shape=[-1],
                                                    dim_range=OrderedDict([
                                                        ('batch_size',
                                                         [bs_range])
                                                    ]))
                # Initialize model
                bert_config = BertConfig(
                    vocab_size=vocab_size,
                    hidden_size=hidden_size,
                    num_hidden_layers=num_layers,
                    num_attention_heads=num_heads,
                    intermediate_size=4 * hidden_size,
                    hidden_act=hidden_act,
                    max_position_embeddings=max_position_embeddings,
                    torch_dtype=torch_dtype,
                )

                output_name = "hidden_states"
                if model == BertModel.__name__:
                    hf_bert = BertModel(
                        bert_config,
                        add_pooling_layer=False).cuda().to(torch_dtype).eval()
                    tensorrt_llm_bert = tensorrt_llm.models.BertModel(
                        num_layers=num_layers,
                        num_heads=num_heads,
                        hidden_size=hidden_size,
                        vocab_size=vocab_size,
                        hidden_act=hidden_act,
                        max_position_embeddings=max_position_embeddings,
                        type_vocab_size=bert_config.type_vocab_size,
                        mapping=tensorrt_llm.Mapping(
                            world_size=world_size,
                            rank=rank,
                            tp_size=world_size),  # TP only
                        dtype=trt_dtype)
                    load_from_hf_bert(tensorrt_llm_bert,
                                      hf_bert,
                                      bert_config,
                                      rank=rank,
                                      tensor_parallel=world_size,
                                      fp16=fp16)
                elif model == BertForQuestionAnswering.__name__:
                    hf_bert = BertForQuestionAnswering(bert_config).cuda().to(
                        torch_dtype).eval()
                    output_name = "logits"
                    tensorrt_llm_bert = tensorrt_llm.models.BertForQuestionAnswering(
                        num_layers=num_layers,
                        num_heads=num_heads,
                        hidden_size=hidden_size,
                        vocab_size=vocab_size,
                        hidden_act=hidden_act,
                        max_position_embeddings=max_position_embeddings,
                        type_vocab_size=bert_config.type_vocab_size,
                        num_labels=
                        2,  # just make it a const here, seems to me not worth as a config
                        mapping=tensorrt_llm.Mapping(
                            world_size=world_size,
                            rank=rank,
                            tp_size=world_size),  # TP only
                        dtype=trt_dtype)
                    load_from_hf_qa_bert(tensorrt_llm_bert,
                                         hf_bert,
                                         bert_config,
                                         rank=rank,
                                         tensor_parallel=world_size,
                                         fp16=fp16)
                else:
                    assert False, f"Unknown model {model}"
                # Prepare
                network.set_named_parameters(
                    tensorrt_llm_bert.named_parameters())

                # Forward
                output = tensorrt_llm_bert(input_ids=input_ids,
                                           input_lengths=input_lengths)

                # Mark outputs
                output_dtype = trt.float16 if fp16 else trt.float32
                output.mark_output(output_name, output_dtype)

            # Build engine
            engine_buffer = builder.build_engine(network, builder_config)
            session = tensorrt_llm.runtime.Session.from_serialized_engine(
                engine_buffer)
            stream = torch.cuda.current_stream().cuda_stream

            # Inference
            # The dtype of input_ids should be queried from the engine,
            # for testing purpose, int32 is fine for now.
            input_ids = torch.randint(100, (batch_size, input_len)).int().cuda()
            input_lengths = input_len * torch.ones(
                (batch_size, ), dtype=torch.int32, device='cuda')

            output_info = session.infer_shapes([
                TensorInfo('input_ids', trt.DataType.INT32,
                           (batch_size, input_len)),
                TensorInfo('input_lengths', trt.DataType.INT32, (batch_size, ))
            ])
            session._print_engine_info()

            outputs = {
                t.name: torch.empty(tuple(t.shape),
                                    dtype=trt_dtype_to_torch(t.dtype),
                                    device='cuda')
                for t in output_info
            }
            assert output_name in outputs, f'{output_name} not found in outputs'
            session.run(inputs={
                'input_ids': input_ids,
                'input_lengths': input_lengths
            },
                        outputs=outputs,
                        stream=stream)
            torch.cuda.synchronize()
            res = outputs[output_name]

            with torch.no_grad():
                hf_outputs = hf_bert.forward(input_ids)
            torch.cuda.synchronize()

            if model == BertModel.__name__:
                ref = hf_outputs.last_hidden_state
                np.testing.assert_allclose(ref.cpu().numpy(),
                                           res.cpu().numpy(),
                                           atol=1e-2,
                                           rtol=1e-2)
            elif model == BertForQuestionAnswering.__name__:
                res_start_logits, res_end_logits = torch.split(res, 1, -1)
                res_start_logits = res_start_logits.squeeze()
                res_end_logits = res_end_logits.squeeze()

                ref_start_logits = hf_outputs.start_logits
                ref_end_logits = hf_outputs.end_logits

                np.testing.assert_allclose(ref_start_logits.cpu().numpy(),
                                           res_start_logits.cpu().numpy(),
                                           atol=1.5e-2)
                np.testing.assert_allclose(ref_end_logits.cpu().numpy(),
                                           res_end_logits.cpu().numpy(),
                                           atol=1.5e-2)

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


# Opus/UrbanSim urban simulation software.
# Copyright (C) 2010-2011 University of California, Berkeley, 2005-2009 University of Washington
# See opus_core/LICENSE

from numpy import sometrue, array, ravel, logical_not, zeros, float32, maximum, equal, int32
from numpy import arcsin, sqrt
from numpy import ma
from scipy.ndimage import maximum_filter, correlate

from opus_core.variables.variable import Variable

from biocomplexity.land_cover.variable_functions import my_attribute_label

class shei(Variable):
    """Shannon's evenness index (SHEI), computed in a 5x5 moving window
      - need {"constants":{"FOOTPRINT":footprint}} in resources when compute
           where footprint=ones(shape=(5,5))"""
    land_cover_type = 'lct'

    def dependencies(self):
        return [my_attribute_label(self.land_cover_type)]

    def _get_cover_types(self, lct):
        """Return a list of landcover types present in the lct grid"""
        x = []
        max_type = int(maximum.reduce(ravel(lct)))
        for itype in range(1, max_type+1):
            if sometrue(ravel(lct) == itype):
                x = x + [itype]
        return array(x)

    def _count_covertypes_within_window(self, lct, cover_type_list, footprint):
        """Return integer array indicating the number of different covertypes
        of interest that are within the moving window (footprint)"""
        m = zeros(shape=lct.shape, dtype=int32)
        for cover_type in cover_type_list:
            m += maximum_filter(equal(lct, cover_type), footprint=footprint)
        return m

    def _compute_pct_cover_type_within_footprint(self, lct, cover_type, footprint, lct_mask):
        """Calculates percentage of one covertype within the footprint"""
        temp = equal(lct, cover_type)
        # use ma.masked_array to prevent divide-by-zero warnings
        mask_invert = logical_not(lct_mask)  # in numpy import ma.masked, 0 == valid
        pixels = ma.masked_array(correlate(mask_invert.astype(int32), footprint, mode="reflect"))
        values = correlate(temp.astype(int32), footprint, mode="reflect")
        return ma.filled(values / pixels.astype(float32), 0)

    def _compute_shannon_evenness_index(self, lct, footprint, lct_mask):
        """Compute Shannon's evenness index"""
        cover_types_list = self._get_cover_types(lct)
        numerator_sum = ma.masked_array(zeros(shape=lct.shape, dtype=float32),
                                           mask=lct_mask)
        for cover_type in cover_types_list:
            pi = self._compute_pct_cover_type_within_footprint(lct, cover_type, footprint, lct_mask)
            numerator_sum += pi*ma.filled(ma.log(pi), 0)

        m = self._count_covertypes_within_window(lct, cover_types_list, footprint)

        return ma.filled(-numerator_sum / ma.log(m), 0).astype(float32)

    def compute(self, dataset_pool):
        constants = dataset_pool.get_dataset('constants')
        footprint = constants["FOOTPRINT"]
        lct = ma.filled(self.get_dataset().get_2d_attribute(self.land_cover_type), 0)
        lct_mask = self.get_dataset().get_mask(is_2d_version=True)
        shei = self._compute_shannon_evenness_index(lct, footprint, lct_mask)
        return self.get_dataset().flatten_by_id(arcsin(sqrt(shei)))

from opus_core.tests import opus_unittest
from opus_core.datasets.dataset_pool import DatasetPool
from opus_core.storage_factory import StorageFactory

from biocomplexity.tests.expected_data_test import ExpectedDataTest

class Tests(ExpectedDataTest):
    variable_name = "biocomplexity.land_cover.shei"

    def test_my_inputs(self):
        storage = StorageFactory().get_storage('dict_storage')

        storage.write_table(
            table_name='land_covers',
            table_data={
                'relative_x': array([1,2,1,2]),
                'relative_y': array([1,1,2,2]),
                "lct": array([1, 2, 1, 4]),
            }
        )

        dataset_pool = DatasetPool(package_order=['biocomplexity'],
                                   storage=storage)

        gridcell = dataset_pool.get_dataset('land_cover')
        gridcell.compute_variables(self.variable_name,
                                   dataset_pool=dataset_pool)
        values = gridcell.get_attribute(self.variable_name)

        should_be = array([1.42948067, 1.18391383, 1.42948103, 1.18391371])

        self.assert_(ma.allclose( values, should_be, rtol=1e-6),
                     msg = "Error in " + self.variable_name)

    def test_on_expected_data(self):
        self.do_test_on_expected_data(["relative_x","relative_y","lct"],
                                      element_atol=0.3)

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


import os
from math import cos, sin, sqrt, pi
import numpy as np
import pyvista as pv
from planegeometry.geometry import Triangle

def Gamma(t, A, nlobes):
	alpha = np.pi/2 - (np.pi/2-A)*np.cos(nlobes*t)
	beta = t + A*np.sin(2*nlobes*t)
	return np.array([
      np.sin(alpha) * np.cos(beta),
      np.sin(alpha) * np.sin(beta),
      np.cos(alpha)
	])

def HopfInverse(p, phi):
	return np.array([
      (1+p[2])*np.cos(phi),
      p[0]*np.sin(phi) - p[1]*np.cos(phi),
      p[0]*np.cos(phi) + p[1]*np.sin(phi),
      (1+p[2])*np.sin(phi)
	]) / np.sqrt(2*(1+p[2]))

def Stereo(q):
	return 2*q[0:3] / (1-q[3])

def F(t, phi, A, nlobes):
	return Stereo(HopfInverse(Gamma(t, A, nlobes), phi))

def HopfTorusMesh(nu=400, nv=200, A=0.44, nlobes=3):
    angle_u = np.linspace(-np.pi, np.pi, nu)
    angle_v = np.linspace(0, np.pi, nv)
    u, v = np.meshgrid(angle_u, angle_v)
    z, x, y = F(u, v, A, nlobes)
    grid = pv.StructuredGrid(x, y, z)
    mesh = grid.extract_geometry().clean(tolerance=1e-6)
    return mesh

# -----------------------------------------------------------------------------
def Hexlet(Center, Radius, HTmesh, nframes, gifname, convert="magick convert", delay=8, bgcolor="#363940", tori_color="orangered"):
    s = Radius # side length of the hexagon
    Coef = 2/3
    a = Coef*(Radius+s/2)/sin(pi/2-2*pi/6)
    I = np.array([a, 0]) # inversion pole
    ## ------------------------------------------------------------------ ////
    O1 = np.array([2*a, 0, 0])
    # interior sphere
    def inversion(M, RRadius):
        II = np.array([Coef*(RRadius+RRadius/2)/sin(pi/2-2*pi/6), 0])
        S = Coef*(RRadius+RRadius/2) * np.array([cos(2*pi/6), sin(2*pi/6)])
        k = np.vdot(S-II, S-II) # negated inversion constant
        M = np.asarray(M, dtype=float)
        IM = M-II
        return II - k/np.vdot(IM,IM)*IM
    SmallRadius = Coef*(Radius-s/2)
    p1 = inversion((SmallRadius,0), Radius)
    p2 = inversion((0,SmallRadius), Radius)
    p3 = inversion((-SmallRadius,0), Radius)
    tr = Triangle(p1, p2, p3)
    cs = tr.circumcircle()

    shift = pi/90
    frames = np.linspace(1, 180, nframes)
    anim = False
    if not gifname is None:
        gif_sansext, file_extension = os.path.splitext(os.path.basename(gifname))
        anim = True
        screenshotfmt = gif_sansext + "_%03d.png"
        screenshotglob = gif_sansext + "_*.png"
    for frame_number in frames:
        pltr = pv.Plotter(window_size=[512,512], off_screen=anim)
        pltr.set_background(bgcolor)
        i = 1
        while i<= 6:
            beta = i*pi/3 - frame_number*shift; # frame from 1 to 180
            ccenter = Coef*Radius*np.array([cos(beta),sin(beta)])
            p1 = inversion((0,Coef*Radius/2)+ccenter,Radius)
            p2 = inversion((Coef*Radius/2,0)+ccenter,Radius)
            p3 = inversion((0,-Coef*Radius/2)+ccenter,Radius)
            tr = Triangle(p1, p2, p3)
            cs = tr.circumcircle()
            center = np.array([cs.center[0], cs.center[1], 0])
            r = cs.radius
            mesh = HTmesh.copy()
            mesh.scale(r*0.05)
            mesh.rotate_z(2*frame_number)
            mesh.translate(center-O1)
            pltr.add_mesh(mesh, color=tori_color, specular=20, smooth_shading=True)
            i += 1
        pltr.set_focus([0,0,0])
        pltr.set_position([2, 0, 8])
        if anim:
            pltr.show(screenshot=screenshotfmt % frame_number)
        else:
            pltr.show()
    if anim:
        os.system(
            convert + (" -dispose previous -loop 0 -delay %d " % delay) + screenshotglob + " " + gifname
        )

Hexlet((0,0,0), 2, HopfTorusMesh(), nframes=1, gifname=None)#"HopfToriSteinerOrbit.gif")


## This module aims to analyze the effect of angle distribuiton
import numpy as np
import math as mt
import matplotlib.pyplot as plt
import Kinematics
import scipy.integrate as integrate
import scipy.optimize as opt

#Particle Property
#kpc
kpc_in_cm = 3.08567758e21
#light speed
vc = 3e10

#Neutrino
M_nu = 0.32 # Unit:ev/c2
E_total_nu = 3.6e53*6.24150913e11 #Total energy of neutrinos #Transfer 2e51erg to unit of MeV
E_per_nu = 10e6 #Mean energy of each neutrino #estimated value

#DM
M_DM = 1e03

#NFW Parameter
rho_s = 0.184e9
rs=24.42*kpc_in_cm

#cross section (Neutrino and DM)
cs = 1e-30

def angle_norm(beta,r,R):
    def f(psi):
        rp = (R**2+r**2-2*R*r*np.cos(psi))**0.5
        cos_theta_psi = ((R**2+(r*np.cos(psi))**2-2*R*r*np.cos(psi))/(R**2+r**2-2*R*r*np.cos(psi)))**0.5

        def g(psi):
            def theta(psi):
                return np.arctan(1/(1/(np.tan(psi))-r/(R*np.sin(psi))))

            theta = theta(psi)
            sec = 1/np.cos(theta)

            return 4*np.tan(theta)*(sec**2)*(1-beta*beta)/((sec**2-beta**2)**2)/np.sin(theta)
        return np.sin(psi)*g(psi)*cos_theta_psi/(rp**2)/2.

    theta_max = np.pi/2
    def psi(theta):
        return theta - np.arcsin(r/R*np.sin(theta))
    psi_max = psi(theta_max)
    return integrate.nquad(f, [[0,psi_max]])[0]

def DM_number(m_dm,e_per_nu,start,end,n_total):
    beta = (E_per_nu**2- M_nu**2)**0.5/(E_per_nu+M_DM)
    gamma = (1-beta**(2))**0.5
    time_delay = np.sum((start-end)**2)**0.5*(1/beta-1)/(vc)

    R = (np.sum((start-end)**2))**0.5
    l = end -start
    def f(r,theta):
        psi = theta - np.arcsin(r/R*np.sin(theta))
        dpsi_dtheta = 1- r/R*np.cos(theta) /(1-(r/R*np.sin(theta))**2)**0.5
        x=(np.sum((start + np.array([r*np.cos(psi),0,r*np.sin(psi)]))**2))**0.5/rs
        cos_theta_psi=((R**2+(r*np.cos(psi))**2-2*R*r*np.cos(psi))/(R**2+r**2-2*R*r*np.cos(psi)))**0.5
        theta = np.arctan(1/(1/(np.tan(psi))-r/(R*np.sin(psi))))
        sec = 1/np.cos(theta)
        dn_domega= (sec**3)*(1-beta*beta)/((sec**2-beta**2)**2)/np.pi

        return np.sin(psi) *dpsi_dtheta *R/(R**2+r**2-2*R*r*np.cos(psi))*dn_domega/(x*(1+x)*(1+x))

    k = n_total*rho_s*cs/m_dm/(4*np.pi)
    r0 = 0.01*rs
    result =integrate.dblquad(f, 0, np.pi/2., lambda theta: r0, lambda theta: R)
    print(result[1])
    L_dm = result[0]*2*np.pi*k/R
    print("DM Number(1/cm^2):"+str(L_dm))
    return L_dm

def DM_flux(m_dm,e_per_nu,start,end,n_total,t):
    beta = (E_per_nu**2- M_nu**2)**0.5/(E_per_nu+M_DM)
    a =(vc)*t
    #print(beta)
    gamma = (1-beta**(2))**0.5
    time_delay = np.sum((start-end)**2)**0.5*(1/beta-1)/(vc)
    R = (np.sum((start-end)**2))**0.5
    l = end -start
    def get_dtheta_dt(r,t):
        def get_theta():
            def f(theta):
                #v = vc *beta *((1-beta**2)/(np.cos(theta)**(-2)-beta**2))**0.5
                eta  = 2* np.arctan(np.tan(theta)/(1-beta**2)**0.5)
                root = ((1+np.cos(eta))**2 + (np.sin(eta)**2) *(1-beta**2))**0.5
                v = vc *beta *root/(1+beta**2*np.cos(eta))
                l = v*(t-r/vc)
                return l**2+ r**2 - R**2 + 2*r*l*np.cos(theta)
            sol = opt.fsolve(f, np.pi/4)
            return sol[0]
        theta = np.abs(get_theta())
        psi = theta - np.arcsin(r/R*np.sin(theta))
        eta  = 2* np.arctan(np.tan(theta)/(1-beta**2)**0.5)
        l = R*np.sin(psi)/np.sin(theta)
        #v = vc *beta *((1-beta**2)/(np.cos(theta)**(-2)-beta**2))**0.5
        eta  = 2* np.arctan(np.tan(theta)/(1-beta**2)**0.5)
        root = ((1+np.cos(eta))**2 + (np.sin(eta)**2) *(1-beta**2))**0.5

        v = vc *beta *root/(1+beta**2*np.cos(eta))
        sec = np.cos(theta)**(-1)
        dl_dtheta = r*np.sin(theta)-r**2*np.sin(2*theta)/((R**2-r**2*np.sin(theta)**2)**0.5)
        #dv_dtheta = -vc*beta*np.sin(eta/2.)/2. *2*sec**2*(1-beta*beta)**0.5/(sec**2-beta**2)
        d_eta_dtheta = 2*sec**2*(1-beta*beta)**0.5/(sec**2-beta**2)

        eba = (1+beta**2*np.cos(eta))
        dv_dtheta = vc*beta *d_eta_dtheta* (  root*(beta**2*np.sin(eta)) - eba/root*(np.sin(eta)+beta**2*np.sin(eta)*np.cos(eta))  )/eba**2

        dt_dtheta = dl_dtheta/v -l/(v**2)*dv_dtheta

        return 1/dt_dtheta, theta, psi

    def f(r):
        dtheta_dt, theta, psi = get_dtheta_dt(r,t)
        dpsi_dtheta = 1- r/R*np.cos(theta) /(1-(r/R*np.sin(theta))**2)**0.5
        x=(np.sum((start + np.array([r*np.cos(psi),0,r*np.sin(psi)]))**2))**0.5/rs
        cos_theta_psi=((R**2+(r*np.cos(psi))**2-2*R*r*np.cos(psi))/(R**2+r**2-2*R*r*np.cos(psi)))**0.5

        sec = 1/np.cos(theta)
        dn_domega= (sec**3)*(1-beta*beta)/((sec**2-beta**2)**2)/np.pi

        return np.sin(psi) *dpsi_dtheta *R/(R**2+r**2-2*R*r*np.cos(psi))*dn_domega/(x*(1+x)*(1+x)) *dtheta_dt

    k = n_total*rho_s*cs/m_dm/(4*np.pi)
    r0 = 0.01*rs
    result = integrate.quad(f,r0,R)[0]
    if result<0:
        return 0
    L_dm = result*2*np.pi*k/R
    return L_dm

def DM_number_original(m_dm,e_per_nu,start,end,n_total):
    beta = (E_per_nu**2- M_nu**2)**0.5/(E_per_nu+M_DM)
    gamma = (1-beta**(2))**0.5
    time_delay = np.sum((start-end)**2)**0.5*(1/beta-1)/(vc)

    R = (np.sum((start-end)**2))**0.5
    l = end -start
    def f(t):
        r=(np.sum((start+l*t)**2))**0.5
        x= r/rs
        return 1/(x*(1+x)*(1+x))
    k = n_total*rho_s*cs/m_dm/(4*np.pi) /R
    r0 = 0.01*rs/R
    L_dm = integrate.nquad(f, [[r0,1.]])[0]*k
    print("DM Number original(1/cm^2):"+str(L_dm))
    return  L_dm

def Spectrum(m_dm,e_per_nu,start,end,n_total):
    beta = (E_per_nu**2- M_nu**2)**0.5/(E_per_nu+M_DM)
    gamma = (1-beta**(2))**0.5
    time_delay = np.sum((start-end)**2)**0.5*(1/beta-1)/(vc)

    R = (np.sum((start-end)**2))**0.5
    l = end -start
    def dL_dEdR(r,theta):

        dE_dtheta = M_DM *(beta**2)*np.sin(theta)/(1-beta**2)
        def f(r):
            psi = theta - np.arcsin(r/R*np.sin(theta))
            x=(np.sum((start + np.array([r*np.cos(psi),0,r*np.sin(psi)]))**2))**0.5/rs
            cos_theta_psi= np.cos(theta-psi)
            sec = 1/np.cos(theta)
            dpsi_dtheta = 1- r/R*np.cos(theta) /(1-(r/R*np.sin(theta))**2)**0.5
            dn_domega= (sec**3)*(1-beta*beta)/((sec**2-beta**2)**2)/np.pi

            return np.sin(psi)*cos_theta_psi*R /(R**2+r**2-2*R*r*np.cos(psi))*dn_domega*dpsi_dtheta/(x*(1+x)*(1+x))
        return  f(r)

    E_0 = M_DM*1.001
    E_f = M_DM/(1-beta**2)*(1+(beta**2))*0.999

    norm = integrate.dblquad(dL_dEdR, 0, np.pi/2, lambda theta: 0, lambda theta: R)[0]*2*np.pi
    print(norm)
    def dL_dE(theta):

        sec = 1/np.cos(theta)
        dE_dtheta = M_DM *(beta**2)/(1-beta**2) *4*(gamma**2)*np.tan(theta)*(sec**2) /(1+(gamma**2)*(np.tan(theta)**2))
        def f(r):
            psi = theta - np.arcsin(r/R*np.sin(theta))
            x=(np.sum((start + np.array([r*np.cos(psi),0,r*np.sin(psi)]))**2))**0.5/rs
            cos_theta_psi= np.cos(theta-psi)
            sec = 1/np.cos(theta)
            dpsi_dtheta = 1- r/R*np.cos(theta) /(1-(r/R*np.sin(theta))**2)**0.5
            dn_domega= (sec**3)*(1-beta*beta)/((sec**2-beta**2)**2)/np.pi

            return np.sin(psi)*cos_theta_psi*R /(R**2+r**2-2*R*r*np.cos(psi))*dn_domega*dpsi_dtheta/(x*(1+x)*(1+x))
        return  integrate.nquad(f, [[0.01*R,R]])[0]  /dE_dtheta

    theta =np.linspace(0.001,np.pi/2,1000)
    E = [M_DM/(1-beta**2)*(1+(beta**2)*(1-(gamma**2)*(np.tan(theta[i])**2)) /(1+(gamma**2)*(np.tan(theta[i])**2)))for i in range(0,1000)]

    spec = [dL_dE(theta[i])/norm for i in range(0,1000)]

    plt.plot(E, spec, color ='blue')
    plt.xlabel('E(eV)')
    plt.ylabel('dL_chi/dE(1/cm**2 eV)')

    plt.show()

if __name__== '__main__':

    beta = (E_per_nu**2- M_nu**2)**0.5/(E_per_nu+M_DM)
    gamma = (1-beta**(2))**0.5
    print(beta)
    print("Angle Norm:"+str(angle_norm(beta,0.5,1)))

    start=np.array([8.7*0.0*kpc_in_cm,0,0.*3.08567758e18])
    end =np.array([8.7*kpc_in_cm,0,0.*3.08567758e18])

    ref = DM_number(M_DM,E_per_nu ,start,end,E_total_nu/E_per_nu)

    t0 = kpc_in_cm/(vc)

    mode = 1
    if mode==0:
        N = 100
        s = np.linspace(-3,-0.5,100)
        #t = [8.7*(0.99+10**(s[i]))*t0 for i in range(len(s))]

        t = np.linspace(8.7*(0.99+9e-3),8.7*(0.99+1e-2),N)*t0
        flux = [DM_flux(M_DM,E_per_nu ,start,end,E_total_nu/E_per_nu,t[i]) for i in range(len(t))]

        t2 = np.linspace(8.7*(0.99+1e-2),8.7*(0.99+1e-1),N)*t0
        flux2 = [DM_flux(M_DM,E_per_nu ,start,end,E_total_nu/E_per_nu,t2[i]) for i in range(len(t))]

        t3 = np.linspace(8.7*(0.99+1e-1),8.7*(0.99+1e-0),N)*t0
        flux3 = [DM_flux(M_DM,E_per_nu ,start,end,E_total_nu/E_per_nu,t3[i]) for i in range(len(t))]
        #ref = DM_number_original(M_DM,E_per_nu ,start,end,E_total_nu/E_per_nu)

        inter = [(flux[i+1]+flux[i])/2*(t[i+1]-t[i])for i in range(len(t)-1)]
        inter2 = [(flux2[i+1]+flux2[i])/2*(t2[i+1]-t2[i])for i in range(len(t)-1)]
        inter3 = [(flux3[i+1]+flux3[i])/2*(t3[i+1]-t3[i])for i in range(len(t)-1)]

        t=np.concatenate((t, t2), axis=None)
        t=np.concatenate((t, t3), axis=None)
        flux=np.concatenate((flux, flux2), axis=None)
        flux=np.concatenate((flux, flux3), axis=None)
        print("ratio:"+str(np.sum(inter+inter2+inter3)/ref))
        plt.plot(t,flux)
        plt.xscale('log')
        plt.yscale('log')
        plt.show()
    elif mode==1:
        s = np.linspace(-8,-4,100)
        r0 = 0.01*rs
        delta_t = 8.7*t0 #- r0/vc
        t = [ delta_t+ 8.7*(10**(s[i]))*t0 for i in range(len(s))]
        flux = [DM_flux(M_DM,E_per_nu ,start,end,E_total_nu/E_per_nu,t[i]) for i in range(len(t))]
        plt.plot(t-delta_t*np.ones(len(t)),flux)
        plt.xscale('log')
        plt.yscale('log')
        plt.xlabel('t(s)')
        plt.ylabel('Flux (#/s/cm^2)')

        plt.show()
        inter = [(flux[i+1]+flux[i])/2*(t[i+1]-t[i])for i in range(len(t)-1)]
        print("ratio:"+str(np.sum(inter)/ref))


from django import forms
from users.models import Project, Project_pictures, Comment, Donation, Reply
import datetime

class AddProjectForm(forms.ModelForm):
    class Meta:
        model = Project
        fields = (
            "title",
            "details",
            "total_target",
            "start_date",
            "end_date",
            "category",
            "tags",
        )
        error_messages = {
            'total_target': {
                'min_value': "Invalid value, target must be greater than zero",
            },
        }
        widgets = {
            "tags": forms.TextInput(attrs={"data-role": "tagsinput", "name": "tags"}),
            "start_date": forms.DateInput(attrs={"type": "date"}),
            "end_date": forms.DateInput(attrs={"type": "date"}),
        }

    def clean(self):
        cleaned_data = super().clean()
        start_date = cleaned_data.get("start_date")
        end_date = cleaned_data.get("end_date")
        if end_date <= start_date:
            msg = "End date should be greater than start date."
            self.add_error("end_date", msg)
        if end_date < datetime.date.today() or end_date == datetime.date.today():
            msg = "End date should be greater than today date."
            self.add_error("end_date", msg)
        if start_date < datetime.date.today():
            msg = "Start date should be greater than today date."
            self.add_error("start_date", msg)

class ImageForm(forms.ModelForm):
    class Meta:
        model = Project_pictures
        fields = ("picture",)
        labels = {"picture": "Images"}
        widgets = {
            "picture": forms.ClearableFileInput(attrs={"multiple": True}),
        }

class CommentForm(forms.ModelForm):
    class Meta:
        model = Comment
        fields = ("content",)
        labels = {"content": ""}
        widgets = {
            "content": forms.TextInput(
                attrs={
                    "class": "comment-input",
                    "placeholder": "New Comment....",
                }
            ),
        }

class ReplyForm(forms.ModelForm):
    class Meta:
        model = Reply
        fields = ("content",)
        labels = {"content": ""}
        widgets = {
            "content": forms.TextInput(
                attrs={
                    "class": "comment-input w-100",
                    "placeholder": "Add reply....",
                }
            ),
        }

class DonateForm(forms.ModelForm):
    class Meta:
        model = Donation
        fields = ("amount",)
        labels = {"amount": ""}
        widgets = {
            "amount": forms.NumberInput(
                attrs={
                    "class": "form-control",
                    "placeholder": "0",
                    "id":"amount"
                }
            ),
        }


import os
os.environ['KERAS_BACKEND'] = 'tensorflow'
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
from konlpy.tag import Okt
from tensorflow.keras.layers import Embedding, Dense, LSTM
from tensorflow.keras.models import Sequential
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import re
import datetime
pos = []
neg = []
posneg = []
stopwords = ['의', '가', '이', '은', '들', '는', '좀', '잘', '걍', '과', '도', '를', '으로', '자', '에', '와', '한', '하다']
# 한글 형태소 분해작업
def run():
    train_data = pd.read_csv("Data/train.csv", encoding='CP949')
    test_data = pd.read_csv("Data/test.csv", encoding='CP949')
    print(train_data.groupby('label').size().reset_index(name='count'))
    print(test_data.groupby('label').size().reset_index(name='count'))
    okt = Okt()
    X_train = []
    for sentence in train_data['headline']:
        temp_X = []
        sentence = re.sub('[-=+,#/\?:^$.@*\"※~&%ㆍ!』\\‘|\(\)\[\]\<\>`\'…\"\“》]', '', str(sentence))
        temp_X = okt.morphs(sentence, stem=True) #토큰화 작업
        temp_X = [word for word in temp_X if not word in stopwords]
        X_train.append(temp_X)
    X_test=[]
    for sentence in test_data['headline']:
        temp_X = []
        sentence = re.sub('[-=+,#/\?:^$.@*\"※~&%ㆍ!』\\‘|\(\)\[\]\<\>`\'…\"\“》]', '', str(sentence))
        temp_X = okt.morphs(sentence, stem=True) #토큰화 작업
        temp_X = [word for word in temp_X if not word in stopwords]
        X_test.append(temp_X)

    # 토큰화된 단어를 정수인코딩
    max_words = 35000
    tokenizer = Tokenizer(num_words=max_words)
    tokenizer.fit_on_texts(X_train)
    X_train = tokenizer.texts_to_sequences(X_train)
    X_test = tokenizer.texts_to_sequences(X_test)
    print("제목의 최대 길이 : ", max(len(l) for l in X_train))
    print("제목의 평균 길이 : ", sum(map(len, X_train))/ len(X_train))
    plt.hist([len(s) for s in X_train], bins=50)
    plt.xlabel('length of Data')
    plt.ylabel('number of Data')
    plt.show()
    # y값 (라벨링)인코딩
    y_train = []
    y_test = []
    # one hot encoding
    for i in range(len(train_data['label'])):
        if train_data['label'].iloc[i] == 1:
            y_train.append([0, 0, 1])
        elif train_data['label'].iloc[i] == 0:
            y_train.append([0, 1, 0])
        elif train_data['label'].iloc[i] == -1:
            y_train.append([1, 0, 0])
    for i in range(len(test_data['label'])):
        if test_data['label'].iloc[i] == 1:
            y_test.append([0, 0, 1])
        elif test_data['label'].iloc[i] == 0:
            y_test.append([0, 1, 0])
        elif test_data['label'].iloc[i] == -1:
            y_test.append([1, 0, 0])

    y_train = np.array(y_train)
    y_test = np.array(y_test)
    # 리스트 셔플
    max_len = 40 #전체 데이터 길이를 30으로 맞춤
    X_train = pad_sequences(X_train, maxlen=max_len)
    X_test = pad_sequences(X_test, maxlen = max_len)

    #긍정, 부정, 중립 3가지 분류를 위한 LSTM, softmax, categorical_crossentropy 적용
    model = Sequential()
    model.add(Embedding(max_words, 100))
    model.add(LSTM(128))
    model.add(Dense(3,activation='softmax'))

    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    history = model.fit(X_train,y_train, epochs=5, batch_size=10, validation_split=0.1)

    predict = model.predict(X_test)
    predict_labels = np.argmax(predict, axis=1)
    original_labels = np.argmax(y_test, axis=1)

    for i in range(30):
        origin_label=""
        if(original_labels[i] == 1):
            origin_label="중립"
        elif (original_labels[i] == 2):
            origin_label = "긍정"
        else:
            origin_label="부정"
        predict_label = ""
        if (predict_labels[i] == 1):
            predict_label = "중립"
        elif (predict_labels[i] == 2):
            predict_label = "긍정"
        else:
            predict_label = "부정"
        print("[", test_data['headline'].iloc[i], "]\t[예측한 라벨 : ",predict_label,"]")
    now = datetime.datetime.now()
    nowDatetime = now.strftime('%Y_%m_%d_%H시%M분%S초')
    model.save('Model'+nowDatetime+'.h5')

if __name__ == '__main__':
    run()


import argparse
import json
import os
import time

# isort: off
import torch
import torch.multiprocessing as mp
import tensorrt as trt
# isort: on
from safetensors import safe_open
from transformers import AutoModelForCausalLM, GPTNeoXConfig
from weight import load_from_hf_gpt_neox

import tensorrt_llm
from tensorrt_llm.builder import Builder
from tensorrt_llm.logger import logger
from tensorrt_llm.mapping import Mapping
from tensorrt_llm.models import quantize_model
from tensorrt_llm.network import net_guard
from tensorrt_llm.plugin.plugin import ContextFMHAType
from tensorrt_llm.quantization import QuantMode

MODEL_NAME = "gptneox"
hf_gpt = None

class StateDict():

    def __init__(self, quant_ckpt_dir):
        self.model_state_dict = safe_open(quant_ckpt_dir,
                                          framework="pt",
                                          device=0)

    def get(self, k):
        return self.model_state_dict.get_tensor(k).cpu()

class GPTQModel():

    def __init__(self, model_dir, quant_ckpt_dir):
        with open(model_dir + '/config.json', 'r') as f:
            model_config = json.load(f)
            self.config = GPTNeoXConfig()
            self.config.vocab_size = model_config['vocab_size']
            self.config.hidden_size = model_config['hidden_size']
            self.config.num_hidden_layers = model_config['num_hidden_layers']
            self.config.num_attention_heads = model_config[
                'num_attention_heads']
            self.config.intermediate_size = model_config['intermediate_size']
            self.config.hidden_act = model_config['hidden_act']
            self.config.rotary_pct = model_config['rotary_pct']
            self.config.rotary_emb_base = model_config['rotary_emb_base']
            self.config.max_position_embeddings = model_config[
                'max_position_embeddings']
            self.config.initializer_range = model_config['initializer_range']
            self.config.layer_norm_eps = model_config['layer_norm_eps']
            self.config.use_cache = model_config['use_cache']
            self.config.bos_token_id = model_config['bos_token_id']
            self.config.eos_token_id = model_config['eos_token_id']
            self.config.tie_word_embeddings = model_config[
                'tie_word_embeddings']
        self.model_state_dict = StateDict(quant_ckpt_dir)

    def state_dict(self):
        return self.model_state_dict

def get_engine_name(model, dtype, tp_size, rank):
    return '{}_{}_tp{}_rank{}.engine'.format(model, dtype, tp_size, rank)

def serialize_engine(engine, path):
    logger.info(f'Serializing engine to {path}...')
    tik = time.time()
    with open(path, 'wb') as f:
        f.write(engine)
    tok = time.time()
    t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
    logger.info(f'Engine serialized. Total time: {t}')

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--world_size',
                        type=int,
                        default=1,
                        help='world size, only support tensor parallelism now')
    parser.add_argument(
        '--model_dir',
        type=str,
        default=None,
        help='The path to HF GPT-NeoX model / checkpoints to read weights from')
    parser.add_argument('--dtype',
                        type=str,
                        default='float16',
                        choices=['float16', 'float32'])
    parser.add_argument(
        '--timing_cache',
        type=str,
        default='model.cache',
        help=
        'The path of to read timing cache from, will be ignored if the file does not exist'
    )
    parser.add_argument('--log_level', type=str, default='info')
    parser.add_argument('--vocab_size', type=int, default=50432)
    parser.add_argument('--n_layer', type=int, default=44)
    parser.add_argument('--n_positions', type=int, default=2048)
    parser.add_argument('--n_embd', type=int, default=6144)
    parser.add_argument('--n_head', type=int, default=64)
    parser.add_argument('--hidden_act', type=str, default='gelu')
    parser.add_argument(
        '--rotary_pct',
        type=float,
        default=0.25,
        help="Percentage of hidden dimensions to allocate to rotary embeddings."
    )
    parser.add_argument('--max_batch_size', type=int, default=64)
    parser.add_argument('--max_input_len', type=int, default=1024)
    parser.add_argument('--max_output_len', type=int, default=1024)
    parser.add_argument('--max_beam_width', type=int, default=1)
    parser.add_argument('--use_gpt_attention_plugin',
                        nargs='?',
                        const='float16',
                        type=str,
                        default=False,
                        choices=['float16', 'float32'])
    parser.add_argument('--use_gemm_plugin',
                        nargs='?',
                        const='float16',
                        type=str,
                        default=False,
                        choices=['float16', 'float32'])
    parser.add_argument('--use_weight_only_quant_matmul_plugin',
                        nargs='?',
                        const='float16',
                        type=str,
                        default=False,
                        choices=['float16'])
    parser.add_argument('--use_weight_only_groupwise_quant_matmul_plugin',
                        nargs='?',
                        const='float16',
                        type=str,
                        default=False,
                        choices=['float16'])
    parser.add_argument(
        '--groupwise_quant_safetensors_path',
        type=str,
        default=None,
        help=
        "The path to groupwise quantized GPT-NeoX model / checkpoints to read weights from."
    )
    parser.add_argument('--use_layernorm_plugin',
                        nargs='?',
                        const='float16',
                        type=str,
                        default=False,
                        choices=['float16', 'float32'])
    parser.add_argument('--parallel_build', default=False, action='store_true')
    parser.add_argument('--enable_context_fmha',
                        default=False,
                        action='store_true')
    parser.add_argument('--enable_context_fmha_fp32_acc',
                        default=False,
                        action='store_true')
    parser.add_argument(
        '--multi_block_mode',
        default=False,
        action='store_true',
        help=
        'Split long kv sequence into multiple blocks (applied to generation MHA kernels). \
                        It is beneifical when batchxnum_heads cannot fully utilize GPU.'
    )
    parser.add_argument('--gpus_per_node', type=int, default=8)
    parser.add_argument(
        '--output_dir',
        type=str,
        default='engine_outputs',
        help=
        'The path to save the serialized engine files, timing cache file and model configs'
    )
    parser.add_argument('--remove_input_padding',
                        default=False,
                        action='store_true')
    parser.add_argument(
        '--use_parallel_embedding',
        action="store_true",
        default=False,
        help=
        'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled'
    )
    parser.add_argument(
        '--embedding_sharding_dim',
        type=int,
        default=1,  # Meta does TP on hidden dim
        choices=[0, 1],
        help=
        'By default the embedding lookup table is sharded along vocab dimension (--embedding_sharding_dim=0). '
        'To shard it along hidden dimension, set --embedding_sharding_dim=1'
        'Note: embedding sharing is only enabled when --embedding_sharding_dim=0'
    )

    parser.add_argument(
        '--strongly_typed',
        default=False,
        action="store_true",
        help=
        'This option is introduced with trt 9.1.0.1+ and will reduce the building time significantly for fp8.'
    )

    args = parser.parse_args()

    logger.set_level(args.log_level)

    if args.model_dir is not None:
        global hf_gpt
        if not args.use_weight_only_groupwise_quant_matmul_plugin:
            logger.info(f'Loading HF GPT-NeoX model from {args.model_dir}...')
            hf_gpt = AutoModelForCausalLM.from_pretrained(args.model_dir)
            args.n_embd = hf_gpt.config.hidden_size
            args.n_head = hf_gpt.config.num_attention_heads
            args.n_layer = hf_gpt.config.num_hidden_layers
            args.n_positions = hf_gpt.config.max_position_embeddings
            args.vocab_size = hf_gpt.config.vocab_size
            args.rotary_pct = hf_gpt.config.rotary_pct
        else:
            assert (
                args.groupwise_quant_safetensors_path is not None
            ), f'Please set the path to the groupwise quantized GPT-NeoX checkpoints with --groupwise_quant_safetensors_path'
            logger.info(
                f'Loading GPTQ quantized HF GPT-NeoX model from {args.groupwise_quant_safetensors_path}...'
            )
            hf_gpt = GPTQModel(args.model_dir,
                               args.groupwise_quant_safetensors_path)
            args.n_embd = hf_gpt.config.hidden_size
            args.n_head = hf_gpt.config.num_attention_heads
            args.n_layer = hf_gpt.config.num_hidden_layers
            args.n_positions = hf_gpt.config.max_position_embeddings
            args.vocab_size = hf_gpt.config.vocab_size
            args.rotary_pct = hf_gpt.config.rotary_pct

    return args

def build_rank_engine(builder: Builder,
                      builder_config: tensorrt_llm.builder.BuilderConfig,
                      engine_name, rank, args):
    '''
       @brief: Build the engine on the given rank.
       @param rank: The rank to build the engine.
       @param args: The cmd line arguments.
       @return: The built engine.
    '''
    kv_dtype = trt.float16 if args.dtype == 'float16' else trt.float32
    rotary_dim = int((args.n_embd // args.n_head) * args.rotary_pct)

    # Initialize Module
    tensorrt_llm_gpt = tensorrt_llm.models.GPTNeoXForCausalLM(
        num_layers=args.n_layer,
        num_heads=args.n_head,
        hidden_size=args.n_embd,
        vocab_size=args.vocab_size,
        hidden_act=args.hidden_act,
        max_position_embeddings=args.n_positions,
        rotary_dim=rotary_dim,
        dtype=kv_dtype,
        mapping=Mapping(world_size=args.world_size,
                        rank=rank,
                        tp_size=args.world_size),  # TP only
        apply_query_key_layer_scaling=builder_config.
        apply_query_key_layer_scaling,
        use_parallel_embedding=args.use_parallel_embedding,
        embedding_sharding_dim=args.embedding_sharding_dim)

    if args.use_weight_only_quant_matmul_plugin or args.use_weight_only_groupwise_quant_matmul_plugin:
        quant_mode = QuantMode.from_description(
            quantize_weights=True,
            quantize_activations=False,
            per_token=False,
            per_channel=False,
            per_group=args.use_weight_only_groupwise_quant_matmul_plugin,
            use_int4_weights=False)
        quantize_kwargs = {}
        if args.use_weight_only_groupwise_quant_matmul_plugin:
            quantize_kwargs = {
                "group_size": 128,
                "zero": True,
            }
        tensorrt_llm_gpt = quantize_model(tensorrt_llm_gpt, quant_mode,
                                          **quantize_kwargs)

    if args.model_dir is not None:
        assert hf_gpt is not None, f'Could not load weights from hf_gpt model as it is not loaded yet.'

        if args.world_size > 1:
            assert (
                args.n_embd % args.world_size == 0
            ), f'Embedding size/hidden size must be divisible by world size.'
            assert (
                args.n_head % args.world_size == 0
            ), f'Number of attention heads must be divisible by world size.'

        load_from_hf_gpt_neox(
            tensorrt_llm_gpt, hf_gpt, (args.dtype == 'float16'), rank,
            args.world_size, args.use_weight_only_groupwise_quant_matmul_plugin)

    # Module -> Network
    network = builder.create_network()
    network.trt_network.name = engine_name
    if args.use_gpt_attention_plugin:
        network.plugin_config.set_gpt_attention_plugin(
            dtype=args.use_gpt_attention_plugin)
    if args.use_gemm_plugin:
        network.plugin_config.set_gemm_plugin(dtype=args.use_gemm_plugin)
    if args.use_layernorm_plugin:
        network.plugin_config.set_layernorm_plugin(
            dtype=args.use_layernorm_plugin)
    assert not (args.enable_context_fmha and args.enable_context_fmha_fp32_acc)
    if args.enable_context_fmha:
        network.plugin_config.set_context_fmha(ContextFMHAType.enabled)
    if args.enable_context_fmha_fp32_acc:
        network.plugin_config.set_context_fmha(
            ContextFMHAType.enabled_with_fp32_acc)
    if args.multi_block_mode:
        network.plugin_config.enable_mmha_multi_block_mode()
    if args.use_weight_only_quant_matmul_plugin:
        network.plugin_config.set_weight_only_quant_matmul_plugin(
            dtype=args.use_weight_only_quant_matmul_plugin)
    if args.use_weight_only_groupwise_quant_matmul_plugin:
        network.plugin_config.set_weight_only_groupwise_quant_matmul_plugin(
            dtype=args.use_weight_only_groupwise_quant_matmul_plugin)
    if args.world_size > 1:
        network.plugin_config.set_nccl_plugin(args.dtype)
    if args.remove_input_padding:
        network.plugin_config.enable_remove_input_padding()
    with net_guard(network):
        # Prepare
        network.set_named_parameters(tensorrt_llm_gpt.named_parameters())

        # Forward
        inputs = tensorrt_llm_gpt.prepare_inputs(args.max_batch_size,
                                                 args.max_input_len,
                                                 args.max_output_len, True,
                                                 args.max_beam_width)
        tensorrt_llm_gpt(*inputs)

    tensorrt_llm.graph_rewriting.optimize(network)

    engine = None

    # Network -> Engine
    engine = builder.build_engine(network, builder_config)
    if rank == 0:
        config_path = os.path.join(args.output_dir, 'config.json')
        builder.save_config(builder_config, config_path)

    return engine

def build(rank, args):
    torch.cuda.set_device(rank % args.gpus_per_node)
    tensorrt_llm.logger.set_level(args.log_level)
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    # when doing serializing build, all ranks share one engine
    apply_query_key_layer_scaling = False
    builder = Builder()

    cache = None
    for cur_rank in range(args.world_size):
        # skip other ranks if parallel_build is enabled
        if args.parallel_build and cur_rank != rank:
            continue
        builder_config = builder.create_builder_config(
            name=MODEL_NAME,
            precision=args.dtype,
            timing_cache=args.timing_cache if cache is None else cache,
            tensor_parallel=args.world_size,  # TP only
            parallel_build=args.parallel_build,
            num_layers=args.n_layer,
            num_heads=args.n_head,
            hidden_size=args.n_embd,
            vocab_size=args.vocab_size,
            hidden_act=args.hidden_act,
            max_position_embeddings=args.n_positions,
            apply_query_key_layer_scaling=apply_query_key_layer_scaling,
            max_batch_size=args.max_batch_size,
            max_beam_width=args.max_beam_width,
            max_input_len=args.max_input_len,
            int8=args.use_weight_only_quant_matmul_plugin
            or args.use_weight_only_groupwise_quant_matmul_plugin,
            max_output_len=args.max_output_len,
            strongly_typed=args.strongly_typed)

        engine_name = get_engine_name(MODEL_NAME, args.dtype, args.world_size,
                                      cur_rank)
        engine = build_rank_engine(builder, builder_config, engine_name,
                                   cur_rank, args)
        assert engine is not None, f'Failed to build engine for rank {cur_rank}'

        if cur_rank == 0:
            # Use in-memory timing cache for multiple builder passes.
            if not args.parallel_build:
                cache = builder_config.trt_builder_config.get_timing_cache()

        serialize_engine(engine, os.path.join(args.output_dir, engine_name))

    if rank == 0:
        ok = builder.save_timing_cache(
            builder_config, os.path.join(args.output_dir, "model.cache"))
        assert ok, "Failed to save timing cache."

if __name__ == '__main__':
    args = parse_arguments()
    tik = time.time()
    if args.parallel_build and args.world_size > 1 and \
            torch.cuda.device_count() >= args.world_size:
        logger.warning(
            f'Parallelly build TensorRT engines. Please make sure that all of the {args.world_size} GPUs are totally free.'
        )
        mp.spawn(build, nprocs=args.world_size, args=(args, ))
    else:
        args.parallel_build = False
        logger.info('Serially build TensorRT engines.')
        build(0, args)

    tok = time.time()
    t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
    logger.info(f'Total time of building all {args.world_size} engines: {t}')


import os.path

import numpy
import pickle
import streamlit as st
import statistics
from matplotlib import pyplot as plt
import seaborn as sb
from PIL import Image

from tools.analyse_tool import CompleteAnalysis
from objects.training_set import TrainingSet
from objects.patient import Patient

def warning():
    color1='#E75919'
    color2='#EE895C'
    color3='#FFFFFF'
    text ='Before starting the analysis, we strongly recommend to load the desired dataset in advance. You can do this in the "Data Loader" tab.'
    st.markdown(
        f'
{text}
',
        unsafe_allow_html=True)

def create_description():
    info_p1 = "This Sepsis Research Analysis focuses on displaying the relation of selected features and " \
              "the occurrence of sepsis. A histogram is used to visualize the collected data."
    st.markdown(info_p1)

SEPSIS_TOOL_CHOICE = {
    "both": "positive + negative",
    "sepsis": "positive",
    "no_sepsis": "negative",
}

def plot_sepsis_analysis(analysis_obj, col2, selected_label, selected_tool):
    min_val = analysis_obj.min_for_label[selected_label][1]
    max_val = analysis_obj.max_for_label[selected_label][1]
    plot_data = analysis_obj.plot_label_to_sepsis[selected_label]
    # getting the min max average to scale the plot proportional
    bins = numpy.linspace(float(min_val), float(max_val),
                          100 if selected_label != 'Gender' else 2)  # removed [1] from min_val
    # Actually Plotting the Histogram
    fig, ax1 = plt.subplots()
    fig.title = "Histogram"  # doesnt work
    if selected_tool == SEPSIS_TOOL_CHOICE["both"]:
        ax1.hist([plot_data[0], plot_data[1]], density=True, color=['r', 'g'], bins=bins, alpha=0.6,
                 label=["Sepsis", "No Sepsis"])
    elif selected_tool == SEPSIS_TOOL_CHOICE["sepsis"]:
        ax1.hist(plot_data[0], bins=bins, alpha=0.6, color="r", label="Sepsis")
    elif selected_tool == SEPSIS_TOOL_CHOICE["no_sepsis"]:
        ax1.hist(plot_data[1], bins=bins, alpha=0.6, color="g", label="No Sepsis")
    ax1.legend()
    col2.pyplot(fig)
    headline = "Further Statistics for the label " + selected_label + ": "
    st.subheader(headline)
    # Displaying further Statistics
    sepsis_mean = round(statistics.mean(plot_data[0]), 5)
    sepsis_median = round(statistics.median(plot_data[0]), 5)
    sepsis_var = round(statistics.variance(plot_data[0]), 5)
    no_sepsis_mean = round(statistics.mean(plot_data[1]), 5)
    no_sepsis_median = round(statistics.median(plot_data[1]), 5)
    no_sepsis_var = round(statistics.variance(plot_data[1]), 5)
    diff_mean = round(sepsis_mean - no_sepsis_mean, 5)
    diff_median = round(sepsis_median - no_sepsis_median, 5)
    diff_var = round(sepsis_var - no_sepsis_var, 5)
    col0, col1, col2, col3 = st.columns(4)
    col0.markdown("**Sepsis**")
    col1.metric("Average", sepsis_mean, diff_mean)
    col2.metric("Median", sepsis_median, diff_median)
    col3.metric("Variance", sepsis_var, diff_var)
    col0, col1, col2, col3 = st.columns(4)
    col0.markdown("**No Sepsis**")
    col1.metric("", no_sepsis_mean)
    col2.metric("", no_sepsis_median)
    col3.metric("", no_sepsis_var)

class SepsisResearch:

    CACHE_CORRELATION_POSTFIX = "frontend-correlation"

    def __init__(self):
        st.markdown("
Histogram for Sepsis Research
",
                    unsafe_allow_html=True)
        create_description()
        col1, col2 = st.columns((1, 2))
        selected_label, selected_set, selected_tool = self.create_selectors(col1)
        analysis_obj, file_name = CompleteAnalysis.get_analysis(selected_label=selected_label,
                                                                selected_tool=selected_tool,
                                                                selected_set=selected_set)
        plot_sepsis_analysis(analysis_obj, col2, selected_label, selected_tool)

        st.markdown(
            "
Correlation of relevant Features
",
            unsafe_allow_html=True)
        col1, col2 = st.columns((1, 2))
        selected_label, selected_set, use_fix_missing_values, use_interpolation\
            = self.__create_correlation_selectors(col1)
        self.__plot_correlations(set=TrainingSet.get_training_set(selected_set), label=selected_label,
                                 fix_missing_values=use_fix_missing_values, use_interpolation=use_interpolation,
                                 col=col2)
        warning()

    def create_selectors(self, col1):
        selected_label = col1.selectbox('Choose a label:', Patient.LABELS)
        selected_set = col1.selectbox('Choose a Set:', TrainingSet.PRESETS.keys())
        selected_sepsis = col1.selectbox('Choose if sepsis positive or negative:',
                                         tuple(SEPSIS_TOOL_CHOICE.values()))
        selected_tool = selected_sepsis
        return selected_label, selected_set, selected_tool

    def __create_correlation_selectors(self, col):
        sepsislabel_indices = [i for i in range(len(Patient.LABELS)) if Patient.LABELS[i] == "SepsisLabel"]
        sepsislabel_index = sepsislabel_indices[0] if len(sepsislabel_indices) > 0 else None
        sample_a_indices = [i for i in range(len(TrainingSet.PRESETS.keys())) if list(TrainingSet.PRESETS.keys())[i] == "rnd Sample A"]
        sample_a_index = sample_a_indices[0] if len(sepsislabel_indices) > 0 else None
        selected_label = col.selectbox("Choose a Label", Patient.LABELS, index=sepsislabel_index, key="corrLabel")          # might be useful to only offer labels that are actually in the selected_set
        selected_set = col.selectbox("Choose a Set:", TrainingSet.PRESETS.keys(), key="corrSet", index=sample_a_index)
        use_fix_missing_values = col.checkbox("Use \"fix missing\"", key="corrFix", value=True)
        use_interpolation = col.checkbox("Use interpolation", key="corrInterpolate", value=True)

        return selected_label, selected_set, use_fix_missing_values, use_interpolation

    def __plot_correlations(self, set: TrainingSet, col, label: str, fix_missing_values: bool, use_interpolation: bool):
        avg_df = set.get_average_df(fix_missing_values=fix_missing_values,
                                    use_interpolation=use_interpolation)
        file_path = set.get_cache_file_path(SepsisResearch.CACHE_CORRELATION_POSTFIX)
        curr_version = "fixed" if fix_missing_values else "no_fixed" + \
                                                          "interpolated" if use_interpolation else "no_interp"
        sorted_corr_df = None
        if not os.path.isfile(file_path):
            print("Frontend SepsisResearch found no cache!")
            d = None
        else:
            print("Frontend SepsisResearch is using cache file:", file_path)
            d = pickle.load(open(file_path, 'rb'))
            if f"sorted_corr_df_{curr_version}" in d.keys():
                sorted_corr_df = d[f'sorted_corr_df_{curr_version}']
            avg_df_corr_without_nan = d['avg_df_corr']
            feature_names = avg_df_corr_without_nan.columns

        if sorted_corr_df is None:  # cache not found or empty
            print("Cache was not found or empty! Calculating ...")
            sepsis_df = set.get_sepsis_label_df()  # no transpose needed
            transposed_df = avg_df.transpose()
            added_sepsis_df = transposed_df
            added_sepsis_df["SepsisLabel"] = sepsis_df.iloc[0:].values
            added_sepsis_df = added_sepsis_df.fillna(0)  # fix NaN problem

            avg_df_corr = added_sepsis_df.corr()
            feature_names = avg_df_corr.columns
            avg_df_corr_without_nan = avg_df_corr.fillna(0)  # if features have no values they are none
            corr_df = avg_df_corr_without_nan[label]
            sorted_corr_df = corr_df.sort_values(ascending=False)
            print("Frontend SepsisResearch writes to cache file:", file_path)
            if d is None:
                pickle.dump({f'sorted_corr_df_{curr_version}': sorted_corr_df, "avg_df_corr": avg_df_corr_without_nan},
                            open(file_path, 'wb'))
            else:
                d[f'sorted_corr_df_{curr_version}'] = sorted_corr_df
                pickle.dump(d, open(file_path, 'wb'))

        # Bar plot of correlation to label
        try:
            fig, ax1 = plt.subplots(1)
            plot_corr_df = sorted_corr_df.drop(label)
            ax1.bar(plot_corr_df.index, plot_corr_df)# = plot_corr_df.plot.bar(x='Features')
            ax1.set_xlabel("Features")
            f = 'fixed values,' if fix_missing_values else ''
            i = 'quadratic interpolation' if use_interpolation else ''
            ax1.set_title(f"Correlation to {label}, {f} {i}")
            props = {"rotation": 90}
            plt.setp(ax1.get_xticklabels(), **props)

            # heat map of feature x feature
            #ax2 = sb.heatmap(data=avg_df_corr_without_nan.to_numpy(), vmin=-1, vmax=1, linewidths=0.5,
            #                 cmap='bwr', yticklabels=feature_names)
            #ax2.set_title(f"Correlations in {set.name}, {f}"
            #              f"{i}")

            # pair plot of greatest features to label
            #important_features = sorted_corr_df.index[:3].tolist()
            #important_features.extend(sorted_corr_df.index[-3:].tolist())
            #selected_labels_df = avg_df.transpose().filter(important_features, axis=1)
            #avg_df_small = selected_labels_df.iloc[:100]  # scatter plot nur 100 patients
            #sb.set_style('darkgrid')
            #pairplot = sb.pairplot(avg_df_small)
            #col.pyplot(pairplot)

            col.pyplot(fig)

        except KeyError:
            info_key_error = "The selected label was not found in the selected dataset. It was probably removed" \
                             " within the imputation. Please select a different label or a different dataset."
            st.markdown(info_key_error)


from fileinput import filename
import logging
import logging
import os
import pandas as pd
from ekorpkit import eKonf

log = logging.getLogger(__name__)

class Quandl:
    def __init__(self, **args):
        from fredapi import Fred
        import nasdaqdatalink

        self.args = eKonf.to_dict(args)
        self.autoload = self.args.get("autoload", True)
        self.name = self.args.get("name")
        self.fred_api_key = self.args.get("fred_api_key")
        self.nasdaq_api_key = self.args.get("nasdaq_api_key")
        self.verbose = self.args.get("verbose", True)
        self.series_id = self.args.get("series_id")
        if isinstance(self.series_id, str):
            self.series_id = [self.series_id]
        elif not isinstance(self.series_id, list):
            self.series_id = []
        self.series_name = self.args.get("series_name")
        self.value_column = self.args.get("value_column") or "value"
        if self.series_name is None:
            if self.series_id:
                self.series_name = "_".join(self.series_id).replace("/", "_")
            else:
                self.series_name = self.value_column

        self.start_date = self.args.get("start_date")
        self.end_date = self.args.get("end_date")
        self.eval_columns = self.args.get("pipeline").get("eval_columns")

        self.output_dir = self.args["output_dir"]
        os.makedirs(self.output_dir, exist_ok=True)
        self.output_file = self.args["output_file"]
        self.force_download = self.args["force_download"]

        self.fred = Fred(api_key=self.fred_api_key)
        nasdaqdatalink.ApiConfig.api_key = self.nasdaq_api_key
        self.data = None

        if self.autoload:
            self.load(
                series_id=self.series_id,
                series_name=self.series_name,
                start_date=self.start_date,
                end_date=self.end_date,
            )

    def get(self, series_id, start_date=None, end_date=None, **kwargs):
        if "/" in series_id:
            return self.get_nasqaq(series_id, start_date, end_date, **kwargs)
        else:
            series = self.get_series(series_id, start_date, end_date, **kwargs)
            df = pd.DataFrame(series, columns=[self.value_column])
            return df

    def get_nasqaq(self, series_id, start_date=None, end_date=None, **kwargs):
        """Return dataframe of requested dataset from Nasdaq Data Link.
        :param series_id: str or list, depending on single dataset usage or multiset usage
                Dataset codes are available on the Nasdaq Data Link website
        :param str api_key: Downloads are limited to 50 unless api_key is specified
        :param str start_date, end_date: Optional datefilers, otherwise entire
            dataset is returned
        :param str collapse: Options are daily, weekly, monthly, quarterly, annual
        :param str transform: options are diff, rdiff, cumul, and normalize
        :param int rows: Number of rows which will be returned
        :param str order: options are asc, desc. Default: `asc`
        :param str returns: specify what format you wish your dataset returned as,
            either `numpy` for a numpy ndarray or `pandas`. Default: `pandas`
        :returns: :class:`pandas.DataFrame` or :class:`numpy.ndarray`
        Note that Pandas expects timeseries data to be sorted ascending for most
        timeseries functionality to work.
        Any other `kwargs` passed to `get` are sent as field/value params to Nasdaq Data Link
        with no interference.
        """
        import nasdaqdatalink

        return nasdaqdatalink.get(
            dataset=series_id, start_date=start_date, end_date=end_date, **kwargs
        )

    def get_series(self, series_id, start_date=None, end_date=None, **kwargs):
        """
        Get data for a Fred series id. This fetches the latest known data, and is equivalent to get_series_latest_release()

        Parameters
        ----------
        series_id : str
            Fred series id such as 'CPIAUCSL'
        start_date : datetime or datetime-like str such as '7/1/2014', optional
            earliest observation date
        end_date : datetime or datetime-like str such as '7/1/2014', optional
            latest observation date
        kwargs : additional parameters
            Any additional parameters supported by FRED. You can see https://api.stlouisfed.org/docs/fred/series_observations.html for the full list

        Returns
        -------
        data : Series
            a Series where each index is the observation date and the value is the data for the Fred series
        """
        return self.fred.get_series(
            series_id,
            observation_start=start_date,
            observation_end=end_date,
            **kwargs,
        )

    def load(
        self,
        series_id,
        series_name=None,
        start_date=None,
        end_date=None,
        filename=None,
        expressions=None,
        index_name="date",
        reset_index=False,
    ):
        if isinstance(series_id, str):
            series_id = [series_id]
        elif not isinstance(series_id, list):
            series_id = []

        if series_name is None:
            series_name = "_".join(series_id).replace("/", "_")
        if filename is None:
            filename = f"{series_name}.parquet"

        if start_date is None:
            start_date = self.start_date
        if end_date is None:
            end_date = self.end_date

        if isinstance(self.eval_columns, dict):
            self.eval_columns["expressions"] = expressions

        if self.verbose:
            print(f"Loading {series_name}{series_id} from {start_date} to {end_date}")

        filepath = os.path.join(self.output_dir, filename)
        if not os.path.exists(filepath) or self.force_download:
            self.data = self._load_series(
                series_id, series_name, start_date, end_date, index_name, reset_index
            )
            eKonf.save_data(self.data, filepath, verbose=self.verbose)
        else:
            log.info(f"{filepath} already exists.")
            self.data = eKonf.load_data(filepath, verbose=self.verbose)
        return self.data.copy()

    def _load_series(
        self,
        series_ids=None,
        series_name=None,
        start_date=None,
        end_date=None,
        index_name="date",
        reset_index=False,
        **kwargs,
    ):

        _dfs = []
        for series_id in series_ids:
            df = self.get(
                series_id=series_id,
                start_date=start_date,
                end_date=end_date,
                **kwargs,
            )
            if len(df.columns) == 1:
                if series_name is None:
                    series_name = df.columns[0]
                df.columns = [self.value_column]
            df = eKonf.pipe(df, self.eval_columns)
            if len(series_ids) > 1:
                df["series_id"] = series_id
            if series_name:
                columns = {
                    col: col.replace(self.value_column, series_name)
                    for col in df.columns
                    if col.startswith(self.value_column)
                }
                df.rename(columns=columns, inplace=True)
            if self.verbose:
                print(df.head())
            _dfs.append(df)

        df = pd.concat(_dfs)
        df.index.name = index_name
        if reset_index:
            df.index.name = series_name + "_" if series_name else "" + df.index.name
            df = df.reset_index()
        return df


import threading
import time
from traceback import print_tb
import numpy as np
import serial.tools.list_ports
import serial
import queue
import PySimpleGUI as sg
import matplotlib.figure as figure
import matplotlib.pyplot as plt

number_of_network_nodes = 3

monitor_serial = 'All' # 'Single' or 'All'

class Artist:
    MAX_SAMPLES = 5000
    def __init__(self, nodes) -> None:
        self.nodes = nodes

        self.rows = 2
        self.cols = 2
        plt.ion()
        self.subplots = [plt.subplots(self.rows, self.cols, num=i, sharex=True) for i in range(len(nodes))]

        self.figures:list[figure.Figure] = [subplot[0] for subplot in self.subplots]
        self.axes: list[tuple[tuple[plt.Axes]]] = [subplot[1] for subplot in self.subplots] # self.rows by self.cols

        self.latest_time = [0] * len(nodes)

        # illuminance
        self.l = [[]] * len(nodes)
        self.l_t = [[]] * len(nodes)
        self.l_lines : list[plt.Line2D] = [ax[0][0].plot([], label='illuminance')[0] for ax in self.axes]

        # reference
        self.r = [[]] * len(nodes)
        self.r_t = [[]] * len(nodes)
        self.r_lines : list[plt.Line2D] = [ax[0][0].plot([], label='reference')[0] for ax in self.axes]

        # duty cycle
        self.DC = [[]] * len(nodes)
        self.DC_t = [[]] * len(nodes)
        self.DC_lines : list[plt.Line2D] = [ax[0][1].plot([], label='duty cycle')[0] for ax in self.axes]

        # integral error
        self.IE = [[]] * len(nodes)
        self.IE_t = [[]] * len(nodes)
        self.IE_lines : list[plt.Line2D] = [ax[1][0].plot([], label='integral error')[0] for ax in self.axes]

        # tracking error
        self.TE = [[]] * len(nodes)
        self.TE_t = [[]] * len(nodes)
        self.TE_lines : list[plt.Line2D] = [ax[1][1].plot([], label='tracking error')[0] for ax in self.axes]

    def setup_figures(self):
        for idx, figure in enumerate(self.figures):
            figure.suptitle(f"Node '{self.nodes[idx]}'")
            self.axes[idx][1][0].set_xlabel('Time (s)')
            self.axes[idx][1][1].set_xlabel('Time (s)')

    def clear_figures(self):
        for idx in range(len(self.nodes)):
            for row in range(self.rows):
                for col in range(self.cols):
                    ax = self.axes[idx][row][col]
                    for line in ax.get_lines():
                        line.set_data([], [])
        self.update_figures()

        # illuminance
        self.l = [[]] * len(self.nodes)
        self.l_t = [[]] * len(self.nodes)

        # reference
        self.r = [[]] * len(self.nodes)
        self.r_t = [[]] * len(self.nodes)

        # duty cycle
        self.DC = [[]] * len(self.nodes)
        self.DC_t = [[]] * len(self.nodes)

        # integral error
        self.IE = [[]] * len(self.nodes)
        self.IE_t = [[]] * len(self.nodes)

        # tracking error
        self.TE = [[]] * len(self.nodes)
        self.TE_t = [[]] * len(self.nodes)

    def update_figures(self):
        for idx in range(len(self.nodes)):
            for row in range(self.rows):
                for col in range(self.cols):
                    ax = self.axes[idx][row][col]
                    ax.set_xlim((self.latest_time[idx] - 5, self.latest_time[idx] + 0.5))
                    ax.legend()
            node_axes = self.axes[idx]
            lux_ax = node_axes[0][0]
            dc_ax = node_axes[0][1]
            ie_ax = node_axes[1][0]
            te_ax = node_axes[1][1]

            dc_ax.set_ylim((-0.05, 1.05))
            if self.l[idx]:
                lux_ax.set_ylim((-2, max(max(max(self.l[idx]), max(self.r[idx])) * 1.1, 10)))
            else:
                lux_ax.set_ylim((-2, 10))
            if self.IE[idx]:
                ie_ax.set_ylim((min(-1, min(self.IE[idx]) - 0.3), max(1, max(self.IE[idx]) + 0.3) ) )
            else:
                ie_ax.set_ylim((-1, 1))
            if self.TE[idx]:
                te_ax.set_ylim((min(-1, min(self.TE[idx]) - 0.3), max(1, max(self.TE[idx]) + 0.3) ) )
            else:
                te_ax.set_ylim((-1, 1))

    @staticmethod
    def wrap(t: list, vals: list):
        vals = vals[-Artist.MAX_SAMPLES:]

        tmin = t[-1] - 6.0
        tmax = t[-1] + 0.5
        t_arr = np.array(t)
        tmin_idx = np.searchsorted(t_arr, tmin)
        tmax_idx = np.searchsorted(t_arr, tmax)
        return t[tmin_idx:tmax_idx], vals[tmin_idx:tmax_idx]

    def func_TE(self, idx, tracking_error, time):
        self.TE_t[idx].append(time)
        self.TE[idx].append(tracking_error)
        self.TE_t[idx], self.TE[idx] = self.wrap(self.TE_t[idx], self.TE[idx])
        self.TE_lines[idx].set_data(self.TE_t[idx], self.TE[idx])

    def func_l(self, idx, illuminance, time):
        self.l_t[idx].append(time)
        self.l[idx].append(illuminance)
        self.l_t[idx], self.l[idx] = self.wrap(self.l_t[idx], self.l[idx])
        self.l_lines[idx].set_data(self.l_t[idx], self.l[idx])

    def func_r(self, idx, reference, time):
        self.r_t[idx].append(time)
        self.r[idx].append(reference)
        self.r_t[idx], self.r[idx] = self.wrap(self.r_t[idx], self.r[idx])
        self.r_lines[idx].set_data(self.r_t[idx], self.r[idx])

    def func_IE(self, idx, integral_error, time):
        self.IE_t[idx].append(time)
        self.IE[idx].append(integral_error)
        self.IE_t[idx], self.IE[idx] = self.wrap(self.IE_t[idx], self.IE[idx])
        self.IE_lines[idx].set_data(self.IE_t[idx], self.IE[idx])

    def func_DC(self, idx, duty_cycle, time):
        self.DC_t[idx].append(time)
        self.DC[idx].append(duty_cycle)
        self.DC_t[idx], self.DC[idx] = self.wrap(self.DC_t[idx], self.DC[idx])
        self.DC_lines[idx].set_data(self.DC_t[idx], self.DC[idx])

ports = serial.tools.list_ports.comports()

serial_ports = []
for port, desc, hwid in sorted(ports):
    print(f"{port}: {desc} [{hwid}]")
for port, desc, hwid in sorted(ports):
    if 'Serial' in desc:
        print(f"Added {port}")
        serial_ports.append(port)

if monitor_serial == 'Single':
    ports_to_monitor = [serial_ports[0]]
else:
    ports_to_monitor = serial_ports

names = [f"Node {i+1}" for i in range(number_of_network_nodes)]
artist = Artist(names)
artist.setup_figures()

commands = (
    ('s t ', 'tracking error', artist.func_TE),
    ('s l ', 'illuminance', artist.func_l),
    ('s r ', 'reference', artist.func_r),
    ('s i ', 'integral error', artist.func_IE),
    ('s d ', 'duty cycle', artist.func_DC),
)

QUEUE_SIZE = 1000

input_height = 5
output_height = 7

inputs_str = [""]*input_height
outputs_str = [[""]*output_height for _ in range(number_of_network_nodes)]

def serial_read(serial: serial.Serial, data_queue_out: queue.Queue, data_queue_in: queue.Queue):
    buf = b''
    while True:
        buf = serial.read_all()
        if buf:
            data_queue_out.put(buf)
        try:
            data_in = data_queue_in.get(block=False)
            if data_in:
                print(f"Writing: {data_in}")
                serial.write(data_in)
                serial.flush()
                data_in = None
        except queue.Empty:
            pass

def update_str_lst(lst, new_str):
    lst[:-1] = lst[1:]
    lst[-1] = new_str

queues_out = [queue.Queue(QUEUE_SIZE) for _ in range(len(ports_to_monitor))]
queues_in = [queue.Queue(QUEUE_SIZE) for _ in range(len(ports_to_monitor))]
serial_objects = []
for port in ports_to_monitor:
    serial_objects += [serial.Serial(port, timeout=None)]
    time.sleep(0.2)

threads = [threading.Thread(target=serial_read, args=(serial_objects[i], queues_out[i], queues_in[i]), daemon=True) for i in range(len(serial_objects))]

[thread.start() for thread in threads]

data_read_buffers = [b'' for _ in range(len(ports_to_monitor))]

sg.theme('Python')
column = [[sg.Text(f'{port}:'), sg.Text(size=(80,output_height), key=port)] for port in ports_to_monitor]

column += [[sg.Text(f'[{ports_to_monitor[0]}] Input:'), sg.Text(size=(80, input_height), key='input_win')]]

TOGGLE_RL_KEYS = [f"{node}_toggle_rl" for node in range(number_of_network_nodes)]
toggle_rl = [sg.Button(f"Node {i+1}", key=TOGGLE_RL_KEYS[i]) for i in range(number_of_network_nodes)]

TOGGLE_TE_KEYS = [f"{node}_toggle_te" for node in range(number_of_network_nodes)]
toggle_te = [sg.Button(f"Node {i+1}", key=TOGGLE_TE_KEYS[i]) for i in range(number_of_network_nodes)]

TOGGLE_IE_KEYS = [f"{node}_toggle_ie" for node in range(number_of_network_nodes)]
toggle_ie = [sg.Button(f"Node {i+1}", key=TOGGLE_IE_KEYS[i]) for i in range(number_of_network_nodes)]

TOGGLE_DC_KEYS = [f"{node}_toggle_dc" for node in range(number_of_network_nodes)]
toggle_dc = [sg.Button(f"Node {i+1}", key=TOGGLE_DC_KEYS[i]) for i in range(number_of_network_nodes)]

layout = [[sg.Column(column, element_justification='l', vertical_alignment='t')],
        [sg.Input(key='SerialInput', do_not_clear=True)],
        [sg.Button('[Enter] Send', bind_return_key=True), sg.Button('Clear Figures')],
        [sg.Text("Toggle Reference/Lux Measurements")] + toggle_rl,
        [sg.Text("Toggle Tracking Error Measurements")] + toggle_te,
        [sg.Text("Toggle Integral Error Measurements")] + toggle_ie,
        [sg.Text("Toggle Duty-Cycle Measurements")] + toggle_dc,
        [sg.Button('Exit')]
]

window = sg.Window('Serial monitor', layout)

_iter = 0
window.read(timeout=1)
while True:
    for idx, buffer_queue in enumerate(queues_out):
        try:
            while not buffer_queue.empty():
                data_read_buffers[idx] += buffer_queue.get(block=False)
            data_string = data_read_buffers[idx].decode('latin-1')
            data_read_buffers[idx] = b''
            if _iter == 0:
                break
            if data_string:
                for line in data_string.split('\n'):
                    if line:
                        update_str_lst(outputs_str[idx], line)
                        for command in commands:
                            if line.startswith(command[0]):
                                args = (line.split(command[0])[1]).split(' ')
                                flargs = [float(arg) for arg in args[1:]]
                                time_of_command = float(args[-1])/1000.0
                                if artist.latest_time[int(args[0])] == 0.0:
                                    artist.latest_time[int(args[0])] = time_of_command
                                if time_of_command < artist.latest_time[int(args[0])] - 1.0: # ignore old samples
                                    continue
                                if time_of_command > artist.latest_time[int(args[0])] + 6000.0: # ignore probably corrupted samples
                                    continue
                                artist.latest_time[int(args[0])] = max(artist.latest_time[int(args[0])], time_of_command)
                                command[2](int(args[0]), *flargs[:-1], time_of_command)

                window[serial_ports[idx]].update('\n'.join(outputs_str[idx]))
        except queue.Empty:
            pass

    _iter += 1
    if _iter % 100 == 0:
        artist.update_figures()
        event, values = window.read(timeout=1)
        if event == sg.TIMEOUT_KEY:
            continue
        if event == sg.WIN_CLOSED or event == 'Exit':
            break
        if event == '[Enter] Send':
            update_str_lst(inputs_str, values['SerialInput'])
            window['input_win'].update('\n'.join(inputs_str))
            queues_in[0].put(f"{values['SerialInput']}\n".encode('latin-1'))
            window['SerialInput'].update('')

        if event == 'Clear Figures':
            artist.clear_figures()

        for idx, toggle_info_key in enumerate(TOGGLE_RL_KEYS):
            if event == toggle_info_key:
                queues_in[0].put(f"s l {idx}\n".encode('latin-1'))
                queues_in[0].put(f"s r {idx}\n".encode('latin-1'))

        for idx, toggle_info_key in enumerate(TOGGLE_IE_KEYS):
            if event == toggle_info_key:
                queues_in[0].put(f"s i {idx}\n".encode('latin-1'))

        for idx, toggle_info_key in enumerate(TOGGLE_TE_KEYS):
            if event == toggle_info_key:
                queues_in[0].put(f"s t {idx}\n".encode('latin-1'))

        for idx, toggle_info_key in enumerate(TOGGLE_DC_KEYS):
            if event == toggle_info_key:
                queues_in[0].put(f"s d {idx}\n".encode('latin-1'))

window.close()


import torch
# import torch.utils
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10
import torch.optim as optim

torch.manual_seed(1234)
ratio = 0.1
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_set = CIFAR10(root="task2_data", train=True,
                    transform=transform, download=True)
test_set = CIFAR10(root="task2_data", train=False,
                   transform=transform, download=True)
train_set, val_set = torch.utils.data.random_split(train_set,
                                                   [int((1 - ratio) * len(train_set)),
                                                    int(ratio * len(train_set))])
test_dataset = torch.utils.data.DataLoader(test_set, batch_size=1, shuffle=False)
train_dataset = torch.utils.data.DataLoader(train_set, batch_size=1, shuffle=False)

class MyModel(torch.nn.Module):
    num_features = 32 * 32 * 3
    num_categories = 10
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(MyModel.num_features, 100)
        self.fc2 = nn.Linear(100, 100)
        self.fc3 = nn.Linear(100, MyModel.num_categories)

        self.optimizer = optim.SGD(self.parameters(), lr=0.001)
        self.criterion = torch.nn.MSELoss()
        # self.loss = 0.0

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def loss_func(self, local_data, local_labels):

        output = self.forward(local_data)
        y = local_labels
        self.optimizer.zero_grad()
        self.loss = self.criterion(output, y)
        return self.loss

    def backward(self, loss):
        self.loss.backward()
        self.optimizer.step()

model = MyModel().to(device)
my_loss = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

num_epochs = 200
for epoch in range(num_epochs):
  sum_loss = 0
  for i, (images, labels) in enumerate(train_dataset):
    images = images.to(device)
    labels = labels.to(device)

    output = model(images)
    loss = my_loss(output, labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    sum_loss += loss.item()

  print('Epoch [%d] Train Loss: %.4f'% (epoch+1, sum_loss/i))

with torch.no_grad():
  correct = total = 0
  for images, labels in test_dataset:
    images = images.to(device)
    labels = labels.to(device)

    output = model(images)
    _, predicted_labels = torch.max(output, 1)
    correct += (predicted_labels == labels).sum()
    total += labels.size(0)
  print('Percent correct: %.3f %%' % ((100 * correct) / (total + 1)))


import json
import requests

FB = 'https://graph.facebook.com'

"""Python client library for the Facebook ads API."""

class AdsAPI(object):
    """A client for the Facebook Ads API."""

    def __init__(self, ad_account_id, access_token=None, expires=None):
        self.ad_account_id = ad_account_id
        self.access_token = access_token
        self.expires = expires
        self.payload = {'access_token': access_token}

    def get_ad_account_info(self, fields):
        url = '%s/%s' % (FB, self.ad_account_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def get_ad_account_users(self):
        url = '%s/%s/users' % (FB, self.ad_account_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def get_broad_targeting_categories(self):
        url = '%s/%s/broadtargetingcategories' % (FB, self.ad_account_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def get_connection_objects(self):
        url = '%s/%s/connectionobjects' % (FB, self.ad_account_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def get_ad_stats(self):
        url = '%s/%s/stats' % (FB, self.ad_account_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def get_ad_campaign_stats(self, campaign_ids=None):
        """Return campaign stats for the given list of campaign ids."""

        payload = {'campaign_ids': json.dumps(campaign_ids)}
        payload.update(self.payload)
        url = '%s/%s/adcampaignstats' % (FB, self.ad_account_id)
        r = requests.get(url, params=payload)
        return r.json()

    def get_ad_group_stats(self, adgroup_ids=None, stats_mode=None):
        payload = {'adgroup_ids': json.dumps(adgroup_ids),
                   'stats_mode': stats_mode}
        payload.update(self.payload)
        url = '%s/%s/adgroupstats' % (FB, self.ad_account_id)
        r = requests.get(url, params=payload)
        return r.json()

    def get_keyword_stats(self, adgroup_id):
        url = '%s/%s/keywordstats' % (FB, adgroup_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def get_conversion_stats(self, campaign_ids=None, adgroup_ids=None):
        url = '%s/%s/adcampaignconversions' % (FB, self.ad_account_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def get_reach_estimate(self, currency, targeting_spec):
        payload = {'currency': currency,
                   'targeting_spec': json.dumps(targeting_spec)}
        payload.update(self.payload)
        url = '%s/%s/reachestimate' % (FB, self.ad_account_id)
        r = requests.get(url, params=payload)
        return r.json()

    def get_targeting_description(self, adgroup_id):
        url = '%s/%s/targetingsentencelines' % (FB, adgroup_id)
        r = requests.get(url, params=self.payload)
        return r.json()

    def search(self, query, type, want_localized_name):
        payload = {'q': query, 'type': type,
                   'want_localized_name': want_localized_name}
        payload.update(self.payload)
        url = '%s/search' % FB
        r = requests.get(url, params=payload)
        return r.json()

    def create_ad_creative(self):
        url = '%s/%s/adcreatives' % (FB, self.ad_account_id)
        r = requests.post(url, params=self.payload)
        return r.json()


import pytest
from gpaw.utilities import compiled_with_libvdwxc

from ase.build import bulk

from gpaw import GPAW, Davidson, Mixer, PW
from gpaw.xc.libvdwxc import vdw_mbeef

from gpaw.test import gen

pytestmark = pytest.mark.skipif(not compiled_with_libvdwxc(),
                                reason='not compiled_with_libvdwxc()')

@pytest.mark.mgga
def test_vdw_libvdwxc_mbeef():
    setup = gen('Si', xcname='PBEsol')

    system = bulk('Si')
    calc = GPAW(mode=PW(200), xc=vdw_mbeef(),
                kpts=(2, 2, 2),
                nbands=4,
                convergence=dict(density=1e-6),
                mixer=Mixer(1.0),
                eigensolver=Davidson(4),
                setups={'Si': setup})
    system.calc = calc
    e = system.get_potential_energy()
    ref = -60.53662105617721
    err = abs(e - ref)
    print('e=%r ref=%r err=%r' % (e, ref, err))
    # It would be reasonable to put 1e-6 as tolerance,
    # but the value changes by 4e-4 depending on libxc version.
    # See https://gitlab.com/gpaw/gpaw/issues/161 .
    assert err < 1e-3, err


from enum import Enum, IntEnum
from typing import List, Optional, Set
from odmantic import Model
import datetime

class EntryType (str, Enum):
    """[summary]
        An entry performed by the user.
    [description]
        A user can record of three types of entries:
        1) Traumatic Event
        2) Journal Entry
        3) Incident
    """
    trauma = "Traumatic Event"
    journal = "Journal"
    incident = "Incident"

class IncidentSeverity (int, Enum):
    """[summary]
        In the case of an incident entry a severity level of the incident.
    [description]
        A user can determine the severity of the incident based on how it affected them.
    """
    uncomfortable = 1
    difficult_to_function = 2
    very_difficult_to_function = 3
    unable_to_function = 4
    required_medical_assistance = 5

class IncidentType (str, Enum):
    """[summary]
        In the case of an incident entry the type of incident.
    [description]
        Identification of what happened during the incident.  Can be more than one.
    """
    panic_attack = "Panic Attack"
    nightmare = "Nightmare"
    excessive_anger_outburst = "Excessive Anger Outburst"
    disassociation  = "Disassociation"
    disconnected = "Disconnected"
    flashback = "Flashback"
    auditory_flashback = "Auditory Flashback"
    intrusive_thoughts = "Intrusive Thoughts"
    hyper_arousal = "Hyper Arousal"
    other = "Other"

class Entry (Model):
    """[summary]
        Base data structure of the PAT API.
    [description]
        All API calls will perform CRUD operations against the record using the appID. No personal identifiable data should be submitted.

    """
    appID : str
    entry_type : EntryType
    entry_date : datetime.date
    journal_entry : Optional[str] = None
    trauma_entry : Optional[str] = None
    incident_severity : Optional[IncidentSeverity]
    incident_type : Optional[Set[IncidentType]] = None
    location : Optional[str] = None
    location_description : Optional[str] = None

    class Config:
        schema_extra = {
            "example": {
                "appID" : "00000001",
                "entry_type" : "Incident",
                "entry_date" : "2020-1-1",
                "journal_entry" : "This is a jouurnal entry which is an optional field and free form text entry",
                "trauma_entry" : "This is a trauma entry which is just like a journal entry, but flagged as a traumatic event to make it easier to share with a healthcare provider.",
                "incident_severity" : 3,
                "incident_type" : "nightmare",
                "location" : "111 Microsoft Way",
                "location_description" : "Work"
            }
        }


"""Module to write data classes information."""

# Python packages
import yaml
import logging as log

# write_data_class ------------------------------------------------------------
def write_data_class(data, file='data.yaml'):
    """Take a data object and write its attributes.

    Args:
        data: object with __dict__ attribute.
        file (:obj:`str`): filename.
    """
    log.debug('Write data')
    with open(file, 'w') as f:
        yaml.dump(data.__dict__, f, default_flow_style=False)
# --------------------------------------------------------------------------- #

# write_solution_as_csv -------------------------------------------------------
def write_solution_as_csv(
        data_dict, header=True, mode='w', file='data.csv', dec=".", sep=";"):
    """Take a data object and write its attributes.

    Args:
        data_dict (:obj:`dict`): dictionary.
        header (:obj:`bool`): if True write the header.
        mode (:obj:`str`): write mode: 'w' or 'a'.
        file (:obj:`str`): filename.
        dec (:obj:`str`): decimal point.
        sep (:obj:`str`): column separator
    """
    with open(file, mode) as f:
        if header:
            f.write(sep.join(data_dict.keys()) + '\n')

        f.write(sep.join([str(v).replace(".", dec)
                          if isinstance(v, (int, float)) else
                          v if v is not None else "None"
                          for v in data_dict.values()]) + '\n')
# --------------------------------------------------------------------------- #