# Copyright (C) 2020  GreenWaves Technologies, SAS

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import math
from functools import reduce
from typing import Sequence

import numpy as np

from graph.types import (FilterParameters, MultiplicativeBiasParameters,
                         Parameters)
from quantization.multiplicative.mult_qtype_base import (MultQTypeBase,
                                                         WrapperMixin)
from quantization.multiplicative.symmetric.mult_mulbias_qtype_new import (
    MultMulBiasQType, MultMulBiasScaleQType)
from quantization.multiplicative.symmetric.symmetric_mult_biases_qtype import \
    SymmetricMultBiasesQType
from quantization.qtype import QType
from quantization.quantization_record_base import (
    ConstantQuantizationRecordBase, InputOutputQuantizationRecordBase,
    QuantizationRecordBase, ScalableFilterQuantizationRecordBase)

class MultQuantizationRecordBase(QuantizationRecordBase):
    def __init__(self, *args, quant_mode="symmetric", info=None, **kwargs):
        super(MultQuantizationRecordBase, self).__init__(*args, info=info, **kwargs)
        if info is None:
            self._info['quant_mode'] = quant_mode
        self._cache = {}

    def clear_cache(self):
        self._cache = {}

    def check_cache(self, name):
        return self._cache.get(name)

    def dequantize_as(self, tensor: np.ndarray, key_name: str, idx: int = None) -> np.ndarray:
        qtype = getattr(self, key_name)
        if idx is not None:
            qtype = qtype[idx]
        return qtype.dequantize(tensor)

    def quantize_as(self, tensor: np.ndarray, key_name: str, idx: int = None) -> np.ndarray:
        qtype = getattr(self, key_name)
        if idx is not None:
            qtype = qtype[idx]
        return qtype.quantize(tensor)

    def dequantize_wrapped(self, tensor: np.ndarray, key_name: str, idx: int = None) -> np.ndarray:
        qtype = getattr(self, key_name)
        if idx is not None:
            qtype = qtype[idx]
        if isinstance(qtype, WrapperMixin):
            return qtype.wrapped.dequantize(tensor)
        return qtype.dequantize(tensor)

    def quantize_wrapped(self, tensor: np.ndarray, key_name: str, idx: int = None) -> np.ndarray:
        qtype = getattr(self, key_name)
        if idx is not None:
            qtype = qtype[idx]
        if isinstance(qtype, WrapperMixin):
            return qtype.wrapped.quantize(tensor)
        return qtype.quantize(tensor)

    def requantize(self, tensor: np.ndarray, key_name: str, idx: int = None) -> np.ndarray:
        qtype = getattr(self, key_name)
        if idx is not None:
            qtype = qtype[idx]

        if isinstance(qtype, WrapperMixin):
            tensor_fp = self.dequantize_wrapped(tensor, key_name, idx=idx)
            tensor_sym = qtype.quantize(tensor_fp)
            return tensor_sym
        if tensor.dtype == np.float32:
            return qtype.quantize(tensor)
        return tensor

    def confirm_dimension(self, out_c_idx: int, key_name: str):
        qtype = getattr(self, key_name)
        qtype.quantized_dimension = out_c_idx

class InputQuantizationMixin(MultQuantizationRecordBase):
    def __init__(self, *args, auto_quantize_inputs=False, auto_dequantize_inputs=False, **kwargs):
        super(InputQuantizationMixin, self).__init__(*args, **kwargs)
        self._auto_quantize_inputs = auto_quantize_inputs
        self._auto_dequantize_inputs = auto_dequantize_inputs

    @property
    def auto_quantize_inputs(self):
        return self._auto_quantize_inputs

    @auto_quantize_inputs.setter
    def auto_quantize_inputs(self, val):
        self._auto_quantize_inputs = val

    @property
    def auto_dequantize_inputs(self):
        return self._auto_dequantize_inputs

    @auto_quantize_inputs.setter
    def auto_quantize_inputs(self, val):
        self._auto_dequantize_inputs = val

    def prepare_inputs(self, params: Parameters,
                       input_tensors: Sequence[np.ndarray], ktype: str = None) -> Sequence[np.ndarray]:
        del params
        if ktype == 'float32' and self.auto_dequantize_inputs:
            return [self.dequantize_wrapped(input_tensor, "in_qs", idx=idx)
                    for idx, input_tensor in enumerate(input_tensors)]
        if ktype == 'symmetric' and self.auto_quantize_inputs:
            return [self.quantize_as(input_tensor, "in_qs", idx=idx) for idx, input_tensor in enumerate(input_tensors)]
        return input_tensors

class OutputQuantizationMixin(MultQuantizationRecordBase):
    def __init__(self, *args, auto_dequantize_outputs=False, auto_quantize_outputs=False, **kwargs):
        super(OutputQuantizationMixin, self).__init__(*args, **kwargs)
        self._auto_quantize_outputs = auto_quantize_outputs
        self._auto_dequantize_outputs = auto_dequantize_outputs

    @property
    def auto_dequantize_outputs(self):
        return self._auto_dequantize_outputs

    @auto_dequantize_outputs.setter
    def auto_dequantize_outputs(self, val):
        self._auto_dequantize_outputs = val

    @property
    def auto_quantize_outputs(self):
        return self._auto_dequantize_outputs

    @auto_quantize_outputs.setter
    def auto_quantize_outputs(self, val):
        self._auto_dequantize_outputs = val

    def get_outputs(self, params: Parameters,
                    output_tensors: Sequence[np.ndarray], ktype: str = None) -> Sequence[np.ndarray]:
        del params
        if ktype == 'symmetric':
            if self._auto_dequantize_outputs:
                return [self.dequantize_as(output_tensor, "out_qs", idx=idx)
                        for idx, output_tensor in enumerate(output_tensors)]
            output_tensors = [self.out_qs[idx].clip(output_tensor)
                              for idx, output_tensor in enumerate(output_tensors)]
        return output_tensors

class MultQuantizationRecord(InputQuantizationMixin, OutputQuantizationMixin, InputOutputQuantizationRecordBase):
    def __init__(self, *args, scale_mul_biases_q=None, info=None, **kwargs):
        super(MultQuantizationRecord, self).__init__(*args, info=info, **kwargs)
        if info is None:
            self.scale_mul_biases_q = scale_mul_biases_q

    @property
    def scale_mul_biases_q(self):
        mul_biases_q = self._info.get('scale_mul_biases_q')
        if mul_biases_q is None:
            mul_biases_q = MultMulBiasScaleQType(dtype=np.uint8)
            self.scale_mul_biases_q = mul_biases_q
        return mul_biases_q

    @scale_mul_biases_q.setter
    def scale_mul_biases_q(self, val):
        self._info['scale_mul_biases_q'] = val

    def set_scale(self, in_idx=0, out_idx=0, extra_scale=1):
        if isinstance(in_idx, int):
            in_scale = self.in_qs[in_idx].scale
        else:
            in_scale = reduce(lambda x, y: x * y, [self.in_qs[idx].scale for idx in in_idx])
        if isinstance(out_idx, int):
            out_scale = self.out_qs[out_idx].scale
        else:
            out_scale = reduce(lambda x, y: x * y, [self.out_qs[idx].scale for idx in out_idx])
        scale_mul_biases_q = self.scale_mul_biases_q
        scale = in_scale * extra_scale / out_scale
        scale_mul_biases_q.scale = scale

class MultAddQuantizationRecord(MultQuantizationRecord):
    def __init__(self, *args, scale_in_mul_biases_q=None, info=None, **kwargs):
        super(MultAddQuantizationRecord, self).__init__(*args, info=info, **kwargs)
        if info is None:
            self._info['scale_in_mul_biases_q'] = scale_in_mul_biases_q

    @property
    def scale_in_mul_biases_q(self):
        mul_biases_q = self._info.get('scale_in_mul_biases_q')
        if mul_biases_q is None:
            mul_biases_q = MultMulBiasScaleQType(dtype=np.uint8)
            self.scale_in_mul_biases_q = mul_biases_q
        return mul_biases_q

    @scale_in_mul_biases_q.setter
    def scale_in_mul_biases_q(self, val):
        self._info['scale_in_mul_biases_q'] = val

    @property
    def scaled_idx(self):
        return 1 if self.in_qs[1].scale > self.in_qs[0].scale else 0

    def set_add_scale(self):
        self.set_scale(in_idx=0 if self.scaled_idx else 1)
        scale_in_mul_biases_q = self.scale_in_mul_biases_q
        scaled_idx = self.scaled_idx
        not_scaled_idx = 0 if scaled_idx else 1
        scale = self.in_qs[scaled_idx].scale / self.in_qs[not_scaled_idx].scale
        scale_in_mul_biases_q.scale = scale

class MultConstantQuantizationRecord(InputQuantizationMixin, InputOutputQuantizationRecordBase,
                                     OutputQuantizationMixin, ConstantQuantizationRecordBase):
    def gen_value(self, value):
        return self.out_qs[0].get_quantized(value, container_is_quantized=self.constants_are_quantized)

class FilterQuantizationMixin(MultQuantizationRecord):
    @property
    def calc_q(self) -> MultQTypeBase:
        return QType(bits=32, q=0, signed=True)

    @property
    def acc_q(self) -> MultQTypeBase:
        return QType(bits=32, q=0, signed=True)

    @property
    def biases_q(self) -> SymmetricMultBiasesQType:
        return self._info.get('biases_q')

    @property
    def weights_q(self) -> MultQTypeBase:
        return self._info.get('weights_q')

    @calc_q.setter
    def calc_q(self, val: MultQTypeBase):
        pass

    @acc_q.setter
    def acc_q(self, val: MultQTypeBase):
        pass

    @biases_q.setter
    def biases_q(self, val: SymmetricMultBiasesQType):
        self._info['biases_q'] = val

    @weights_q.setter
    def weights_q(self, val: MultQTypeBase):
        self._info['weights_q'] = val

    @staticmethod
    def rescale(arr, from_scale, to_scale):
        return np.floor((arr * from_scale/to_scale) + 0.5).astype(arr.dtype)

class MultScalableFilterQuantizationRecord(FilterQuantizationMixin, ScalableFilterQuantizationRecordBase):
    def __init__(self, *args,
                 weights_q: MultQTypeBase = None,
                 biases_q: SymmetricMultBiasesQType = None,
                 mul_biases_q: Sequence[MultMulBiasQType] = None,
                 calc_q: QType = None,
                 acc_q: QType = None,
                 enable_prenorm=False,
                 info=None,
                 **kwargs):
        super(MultScalableFilterQuantizationRecord, self).__init__(*args, info=info, **kwargs)
        if info is None:
            self._info['calc_q'] = calc_q
            self._info['acc_q'] = acc_q
            self._info['biases_q'] = biases_q
            self._info['weights_q'] = weights_q
            self._info['mul_biases_q'] = mul_biases_q
            self._info['enable_prenorm'] = enable_prenorm
        self.biases_q.link(self.weights_q, self.in_qs[0])

    @property
    def unwrap(self):
        return self._unwrap

    @unwrap.setter
    def unwrap(self, val):
        self._unwrap = val
        self.biases_q.link(self.weights_q, self.in_qs[0])

    def compute_prenorm(self, params: FilterParameters):
        if not self.enable_prenorm:
            return 0
        max_bits = self.in_qs[0].bits - 1 + self.weights_q.bits - 1 + 1 + \
            math.ceil(math.log2(params.filter.in_c * params.filter.h * params.filter.w))
        spare_bits = 31 - max_bits
        if self.mul_biases_q.dtype == np.int8:
            bits = 7
        elif self.mul_biases_q.dtype == np.uint8:
            bits = 8
        return max(0, bits - spare_bits)

    @property
    def enable_prenorm(self) -> bool:
        return self._info.get('enable_prenorm')

    @enable_prenorm.setter
    def enable_prenorm(self, val: bool):
        self._info['enable_prenorm'] = val

    @property
    def mul_biases_q(self) -> MultMulBiasQType:
        return self._info.get('mul_biases_q')

    @mul_biases_q.setter
    def mul_biases_q(self, val: MultMulBiasQType):
        self._info['mul_biases_q'] = val

    def get_quantized_bias_offset(self, params, weights):
        # input zero correction is sum(W * Zin) by out_c if weights are channel scaled
        axis = tuple([idx for idx in range(4) if idx != params.filter.get_order_idx('out_c')])
        return np.sum(np.multiply(self.in_qs[0].zero_point,
                                  weights,
                                  dtype=np.int32),
                      dtype=np.int32,
                      axis=axis)

    @property
    def biases_zero_correction(self):
        # output zero correction is So/(Si * Sw) * ZPo by out_c if weights are channel scaled
        scale = self.out_qs[0].scale / (self.in_qs[0].scale * self.weights_q.scale)
        return np.floor((self.out_qs[0].zero_point * scale) + 0.5).astype(np.int32)

    def prepare_biases(self, params: Parameters, biases: np.ndarray,
                       weights: np.ndarray, ktype: str = None) -> np.ndarray:
        if ktype == 'float32':
            return self.biases_q.get_dequantized(biases,
                                                 container_is_quantized=self.constants_are_quantized).astype(np.float32)
        if ktype == 'symmetric':
            return self.gen_biases(params, biases, weights)
        raise ValueError()

    def prepare_weights(self, params: Parameters,
                        weights: np.ndarray, ktype: str = None) -> np.ndarray:
        self.confirm_dimension(params.filter.get_order_idx('out_c'), 'weights_q')
        if ktype == 'float32':
            weights = self.weights_q.get_dequantized(weights,
                                                     container_is_quantized=self.constants_are_quantized)
            return weights.astype(np.float32)
        if ktype == 'symmetric':
            return self.gen_weights(params, weights)
        raise ValueError()

    def gen_weights(self, params: Parameters, weights: np.ndarray) -> np.ndarray:
        return self.weights_q.get_quantized(weights,
                                            container_is_quantized=self.constants_are_quantized)

    def gen_biases(self, params: Parameters, biases: np.ndarray, weights: np.ndarray) -> np.ndarray:
        biases = self.biases_q.get_quantized(
            biases, container_is_quantized=self.constants_are_quantized)
        if self.in_qs[0].zero_point != 0:
            biases -= self.get_quantized_bias_offset(params,
                                                     self.requantize(weights, 'weights_q'))
        if self.out_qs[0].zero_point != 0:
            biases += self.biases_zero_correction
        return biases

    def gen_mul_biases(self, params: MultiplicativeBiasParameters) -> np.ndarray:
        if isinstance(self.mul_biases_q, MultMulBiasQType):
            self.mul_biases_q.pre_normalization = self.compute_prenorm(params)
        return self.mul_biases_q.qbiases

    def apply_multiplicative_bias(self, params: FilterParameters, input_tensor: np.ndarray,
                                  axis: int, ktype: str = None):
        if ktype == 'float32':
            return input_tensor
        if ktype == 'symmetric':
            if isinstance(self.mul_biases_q, MultMulBiasQType):
                self.mul_biases_q.pre_normalization = self.compute_prenorm(params)
            input_tensor = self.mul_biases_q.apply_scales(input_tensor, axis)
            return input_tensor.astype(np.int32)


import numpy as np, os
from astropy.io import fits,ascii
from astropy.table import Table,vstack,Column
from scipy.ndimage.filters import gaussian_filter1d
from multiprocessing import Pool

#This is a version of evo_tracks.py for this project, just meant to grab some example tracks.
#It's been super messed with, so it only works for a few tracks from BPASS v2.2.
#The original is in Conti/code

bpass_dir = '/Volumes/shoobert/Research/UW/Diagnostics/data/BPASS/BPASSv2.2/example_tracks/'

#A dict to help us access the BPASS tracks
bpass_col_names = ['Step','Age','logR_1','logTe_1','logL_1','M_1','M_He_core_1','M_CO_core_1',
            'M_ONe_core_1','NULL','X','Y','C','N','O','Ne','M_H_star','M_He_star','M_C_star',
            'M_N_star','M_O_star','M_Ne_star','M_Mg_star','M_Si_star','M_Fe_star','E_bind_env',
            'E_bind_star','M_rem_weakSN','M_eject_weakSN','M_rem_SN','M_eject_SN',
            'M_rem_superSN','M_eject_superSN','L_ang_bin','P_bin','loga','M_1_alt','M_2',
            'M_tot','DMW_1','DMW_2','DMA_1','DMA_2','DMR_1','DMR_2','L_ang_dot','logR_2',
            'logTe_2','logL_2','RLOF_2','IMF_sys','IMF_rej','V-I','U','B','V','R','I','J','H',
            'K','u','g','r','i','z','f300w','f336w','f435w','f450w','f555w','f606w','f814w',
            'U_2','B_2','V_2','R_2','I_2','J_2','H_2','K_2','u_2','g_2','r_2','i_2','z_2',
            'f300w_2','f336w_2','f435w_2','f450w_2','f555w_2','f606w_2','f814w_2','logQ',
            'logFUV','logNUV']

bpass_col_dict = {col:i+1 for i,col in enumerate(bpass_col_names)}
bpass_index_dict = {i+1:col for i,col in enumerate(bpass_col_names)}

default_bpass_col_names = ['Age','logR_1','logTe_1','logL_1','M_1','X','Y','C','N','O','P_bin','M_2','DMW_1','DMW_2','DMA_1','DMA_2','DMR_1','DMR_2','logR_2','logTe_2','logL_2','U','B','V','R','I','J','H','K','u','g','r','i','z','f300w','f336w','f435w','f450w','f555w','f606w','f814w','U_2','B_2','V_2','R_2','I_2','J_2','H_2','K_2','u_2','g_2','r_2','i_2','z_2','f300w_2','f336w_2','f435w_2','f450w_2','f555w_2','f606w_2','f814w_2']

bpass_ms = np.array([300,200,150,120,100,80,70,60,50,40,35,30,25,24,23,22,21,20,19,18,17,16,15,14,13,12,11,10,9.5,9,8.5,8,7.5,7,6.5,6,5.5,5,4.5,4,3.7,3.5,3.2,3,2.7,2.5,2.3,2.1,2,1.9,1.8,1.7,1.6,1.5,1.4,1.3,1.2,1.1,1,0.9,0.8,0.6,0.5,0.4,0.3,0.2,0.1])
bpass_logPs = np.arange(0,4.2,0.2)
bpass_qs = np.arange(0.1,1,0.1)

class track:
    """#Object that contains, at minimum, an initial mass, an array of times, current mass
    luminosity, temperature, as well as any other info. Takes an astropy table or 'BPASS'
    as the track keyword. If BPASS tracks are enabled, the bpass parameters are used to
    search for the right track and read it in. If you want parameters that aren't in the above
    list, give bpass_cols = [LIST OF COLUMN NAMES]. Make sure the names are from the above
    list of available column names"""
    def __init__(self, mini = 9.0, mini_keyword = 'Mini', tracks = 'BPASS', Z = 'z014', bpass_M = 1.0, bpass_q = 0.5, bpass_logP = 2, bpass_cols = 'default', **kwargs):
        self.mini = mini
        self.kwargs = kwargs
        self.Z = Z
        assert tracks=='BPASS', "This version of evo_tracks is only for demo purposes"
        if bpass_M % 1 == 0:
            m_str = int(bpass_M)
        else:
            m_str = bpass_M
        q_str = bpass_q
        if not isinstance(bpass_logP, str):
            if bpass_logP % 1 == 0:
                P_str = int(bpass_logP)
        else:
            P_str = bpass_logP

        #We're only working with solar metallicity for now
        if not ((self.Z == 'z014')|(self.Z == 'z002')):
            raise NotImplementedError("Unsupported metallicity")

        #Check to see if user wants more columns than the default
        names = default_bpass_col_names
        if bpass_cols != 'default':
            if type(bpass_cols) != list:
                raise TypeError("This argument should be a list of additional columns to include")
            else:
                #For each name in the list, check if its already included, then include it
                for new_name in bpass_cols:
                    if new_name in names:
                        pass
                    else:
                        names.append(new_name)

        #Get the indices of the columns for reading in
        index = [bpass_col_dict[col] for col in names]
        if bpass_logP == 'inf':
            file_name = bpass_dir+'sneplot-{0}-{1}'.format(Z,str(m_str))
        else:
            file_name = bpass_dir+'/sneplot-{0}-{1}-{2}-{3}'.format(Z,str(m_str),str(q_str),str(P_str))

        #Check if the track exists in the directory...
        if os.path.isfile(file_name):

            #read in only given columns
            track = ascii.read(file_name,include_names = ['col{}'.format(i) for i in index])

            #rename columns for easy access later
            for i,name in zip(index,names):
                track['col{}'.format(i)].name = name

            #Create columns with info to remember the initial mass, mass ratio, period
            M1_i = Column([bpass_M for i in range(len(track))],name = 'M1_i')
            q_i = Column([bpass_q for i in range(len(track))],name = 'q_i')
            logP_i = Column([bpass_logP for i in range(len(track))],name = 'logP_i')

            track.add_columns([M1_i,q_i,logP_i])

            #Assign object variable with this track
            self.track = track

        else:
            print(file_name)
            raise FileNotFoundError('That combo of mass, q and logP does not exist')

    def get_time(self, time, time_keyword = 'Time'):
        #interpolate the values in the track to a given time. If the time is invalid,
        #just return 0 (if too early) or nan (if it's exploded)
        #This is disgusting and hacky as all hell.
        #I'm so sorry future me.
        out_arr = []
        out_types = []
        for col in self.track.colnames:
            if (self.track[col].dtype == '<U1')|(self.track[col].dtype == '<U13'): #if the column is a string and we can't interpolate values, just give me the value from the closest time step
                try:
                    close_idx = np.argmin(np.abs(self.track[time_keyword].data - time))
                    val = self.track[col].data[close_idx]
                except ValueError:
                    val = []
                    for t in time:
                        close_idx = np.argmin(np.abs(self.track[time_keyword].data - t))
                        val.append(self.track[col].data[close_idx])

            else:
                val = np.interp(time, self.track[time_keyword].data, self.track[col].data, left = 0.0, right = np.nan)

            out_arr.append(np.array(val))

            if self.track[col].dtype == '<U13':
                out_types.append(float)
            else:
                out_types.append(self.track[col].dtype)

        return Table(np.array(out_arr).T, names = self.track.colnames, dtype = out_types)


import numpy as np
import pandas as pd
from collections import Counter

## read mastersheet

#df = pd.read_excel('mastersheet_reformatted.xlsx', sheet_name='Dataset_full')
df = pd.read_excel('mastersheet_combined_reformatted.xlsx', sheet_name='Dataset_full')
# rename columns
df = df.rename(columns={k:k.strip() for k in df.columns})
df = df.rename(columns={'Total Points (Add age score to comorbidy score)' : 'CCI'})
# exclude two patients without EEG
df = df[(df.SID!='AMSD086') & (df.SID!='AMSD153')].reset_index(drop=True)
df['Eval. date/time'] = pd.to_datetime(df['Eval. date/time'])

## generate sids

df['MRN'] = df.MRN.astype(str)
two_mrn_ids = np.where(df.MRN.str.contains('\('))[0]
df.loc[two_mrn_ids, 'MRN'] = df.MRN.iloc[two_mrn_ids].str.replace('(','_').str.replace(')','').str.split('_', expand=True)[1]
underscore_ids = np.where(df.MRN.str.contains('_'))[0]
df.loc[underscore_ids, 'MRN'] = df.MRN.iloc[underscore_ids].str.replace('_','')
assert np.all(df.MRN.str.isnumeric())
sid_names = ['SID', 'MRN']

## define eeg names

eeg_names = [
#'PDR (Posterior dominant rhythm) (>=8 Hz); If present - specify highest freq.',
'Sleep patterns (Spindles, K-complex, Vertex waves)',
'Symmetry (e.g. no focal slowing)',

'Generalized/Diffuse delta slowing',
'Generalized/Diffuse theta slowing',
#'Diffuse slowing - Either theta or delta or GRDA',
#'Excess/Diffuse alpha',
#'Excess/Diffuse beta',
'Focal/Unilateral delta slowing',
'Focal/Unilateral theta slowing',
#'Focal slowing - Either theta or delta or LRDA',
'GRDA (Generalized rhythmic delta activity) (= FIRDA - frontal intermittent rhythmic delta activity)',
'LRDA (Lateralized rhythmic delta activity)',
'Extreme delta brush', ##

#'Periodic discharges - LPD or GPD or BiPD',
#'Any IIIC: LPD or GPD or BiPD or LRDA or Sz or NCSE or TPW (TPW or GPD with TP morphology)',
'LPD (Lateralized periodic discharges) (=PLED - Periodic lateralized epileptiform discharges)',

'GPD (Generalized periodic discharges) (=GPED/PED) (Not triphasic)',
'GPD with Triphasic morphology',
'Triphasic waves',

#'Sporadic epileptiform discharges (=sporadic discharges)',
'BIPD (bilateral indep. periodic discharges) (=BIPLED - Bilateral independent periodic lateralized epileptiform discharges)', ##
'BIRDs (brief potentially ictal rhythmic discharges)',

'Discrete seizures: generalized',
#'Discrete seizures: focal',
'Non convulsive status epilepticus: generalized', ##
'Non convulsive status epilepticus: focal',

'Burst suppression with epileptiform activity', ##
'Burst suppression without epileptiform activity', ##
'Intermittent brief attenuation',
'Moderately low voltage',
'Extremely low voltage / electrocerebral silence', ##
'EEG Unreactive',] ##

# remove [] in EEG features
for col in eeg_names:
    bracket_ids = df[col].astype(str).str.contains('\[')
    if np.any(bracket_ids):
        df.loc[bracket_ids, col] = df[col][bracket_ids].str.replace(']','').str.split('[',expand=True)[0].str.replace('+','').astype(float)

# GPDnTPW/GPDTPW/TPW --> GPD
eeg_names.remove('GPD (Generalized periodic discharges) (=GPED/PED) (Not triphasic)')
eeg_names.remove('GPD with Triphasic morphology')
eeg_names.remove('Triphasic waves')
eeg_names.append('GPD')
df['GPD'] = ((df['GPD (Generalized periodic discharges) (=GPED/PED) (Not triphasic)'] + df['GPD with Triphasic morphology'] + df['Triphasic waves'])>0).astype(float)
df = df.drop(columns=['GPD (Generalized periodic discharges) (=GPED/PED) (Not triphasic)'])
df = df.drop(columns=['GPD with Triphasic morphology'])
df = df.drop(columns=['Triphasic waves'])

"""
# combine G delta slowing, G theta slowing, and GRDA
eeg_names.remove('Generalized/Diffuse delta slowing')
eeg_names.remove('Generalized/Diffuse theta slowing')
eeg_names.remove('GRDA (Generalized rhythmic delta activity) (= FIRDA - frontal intermittent rhythmic delta activity)')
eeg_names.append('Generalized/Diffuse delta or theta slowing or GRDA')
df['Generalized/Diffuse delta or theta slowing or GRDA'] = ((df['Generalized/Diffuse delta slowing'] + df['Generalized/Diffuse delta slowing'] + df['GRDA (Generalized rhythmic delta activity) (= FIRDA - frontal intermittent rhythmic delta activity)'])>0).astype(float)
df = df.drop(columns=['Generalized/Diffuse delta slowing'])
df = df.drop(columns=['Generalized/Diffuse theta slowing'])
df = df.drop(columns=['GRDA (Generalized rhythmic delta activity) (= FIRDA - frontal intermittent rhythmic delta activity)'])
"""

# combine BIPD and GPD
eeg_names.remove('GPD')
eeg_names.remove('BIPD (bilateral indep. periodic discharges) (=BIPLED - Bilateral independent periodic lateralized epileptiform discharges)')
eeg_names.append('GPD or BIPD')
df['GPD or BIPD'] = ((df['GPD'] + df['BIPD (bilateral indep. periodic discharges) (=BIPLED - Bilateral independent periodic lateralized epileptiform discharges)'])>0).astype(float)
df = df.drop(columns=['GPD', 'BIPD (bilateral indep. periodic discharges) (=BIPLED - Bilateral independent periodic lateralized epileptiform discharges)'])

# reverse normal EEG features
normal_EEG_names = [
    #'PDR (Posterior dominant rhythm) (>=8 Hz); If present - specify highest freq.',
    'Sleep patterns (Spindles, K-complex, Vertex waves)',
    'Symmetry (e.g. no focal slowing)']
for col in normal_EEG_names:
    if col in df.columns:
        df.loc[:,col] = 1-df[col]
        eeg_names.remove(col)
        eeg_names.append('no '+col)
df = df.rename(columns={col:'no '+col for col in normal_EEG_names if col in df.columns})

# combine focal delta and theta and asymmetry
df['Asymmetry'] = ((df['no Symmetry (e.g. no focal slowing)'] + df['Focal/Unilateral delta slowing'] + df['Focal/Unilateral theta slowing'])>0).astype(int)
df = df.drop(columns=['no Symmetry (e.g. no focal slowing)', 'Focal/Unilateral delta slowing', 'Focal/Unilateral theta slowing'])
eeg_names.append('Asymmetry')
eeg_names.remove('no Symmetry (e.g. no focal slowing)')
eeg_names.remove('Focal/Unilateral delta slowing')
eeg_names.remove('Focal/Unilateral theta slowing')

worst_delirium_names = [
'Extreme delta brush',
#'BIPD (bilateral indep. periodic discharges) (=BIPLED - Bilateral independent periodic lateralized epileptiform discharges)',
'Non convulsive status epilepticus: generalized',
'Extremely low voltage / electrocerebral silence',
'Burst suppression with epileptiform activity',
'Burst suppression without epileptiform activity',
'EEG Unreactive',
]

## generate y names (label)

ynames = ['CAM-ICU (0/1)', 'CAM-S SF (0-7)', 'CAM-S LF (0-19)', 'Prorated LF Score', '3D-CAM (0/1)', '3D-CAM-S SF (0-7)', 'LOS (Days)', 'Deceased at hosp disch (0=N; 1=Y)', 'Deceased at 3-mo post disch.  (0=N, 1=Y, 2=Unk)', 'Disch. GOS', 'Disch. GOSE']

col = 'Prorated LF Score'
ids = df[col].astype(str).str.contains('/')
tmp = df[col][ids].str.split('/',expand=True)
tmp0 = tmp[0].str.replace('_','').values.astype(float)
tmp1 = tmp[1].str.replace('_','').values.astype(float)
df.loc[ids, col] = tmp0/tmp1

col = 'Deceased at 3-mo post disch.  (0=N, 1=Y, 2=Unk)'
ids = df[col].astype(str).str.contains('\(likely')
tmp = df[col][ids].str.replace(')','').str.split('\(likely',expand=True)
df.loc[ids, col] = tmp[1].astype(float)

## generate info

info_names = ['Age', 'Gender', 'Pre-hosp. GOS', 'Pre-hosp. GOSE', 'Type of EEG (routine, LTM)', 'Eval. within epoch (Y=1, N=0)', 'Eval. date/time']

df['Gender'] = df.Gender.astype(str).str.strip().str.upper()

col = 'Pre-hosp. GOS'
ids = df[col].astype(str)=='_'
df.loc[ids, col] = np.nan

col = 'Pre-hosp. GOSE'
ids = (df[col].astype(str)=='_') | (df[col].astype(str)=='?')
df.loc[ids, col] = np.nan

col = 'Type of EEG (routine, LTM)'
df[col] = df[col].str.strip().str.upper()
df.loc[df[col]=='LTM',col] = 1
df.loc[df[col]=='ROUTINE',col] = 0

col = 'Eval. within epoch (Y=1, N=0)'
ids = df[col].astype(str).str.contains('\?')
df.loc[ids, col] = df[col][ids].str.replace('?','')

counts = [df[x].sum() for x in eeg_names]
df_count = pd.DataFrame(data={'EEGName':eeg_names, 'Count':counts})
df_count = df_count.sort_values('Count', ascending=False).reset_index(drop=True)

# save
import pdb;pdb.set_trace()
with pd.ExcelWriter('data_to_fit_not_combine_slowing_Aaron.xlsx') as writer:
    df[sid_names+eeg_names].to_excel(writer, sheet_name='X', index=False)
    df[sid_names+ynames].to_excel(writer, sheet_name='y', index=False)
    df[sid_names+info_names].to_excel(writer, sheet_name='info', index=False)

    pd.DataFrame(data={'EEGName':worst_delirium_names}).to_excel(writer, sheet_name='worst_delirium_names', index=False)
    df_count.to_excel(writer, sheet_name='counts', index=False)


import logging
import numpy as np
from os import path
from typing import Optional, List

import autofit as af
import autoarray as aa

logger = logging.getLogger(__name__)

logger.setLevel(level="INFO")

class Preloads(aa.Preloads):
    def __init__(
        self,
        w_tilde: Optional[aa.WTildeImaging] = None,
        use_w_tilde: Optional[bool] = None,
        blurred_image: Optional[aa.Array2D] = None,
        traced_grids_of_planes_for_inversion: Optional[aa.Grid2D] = None,
        sparse_image_plane_grid_list_of_planes: Optional[aa.Grid2D] = None,
        relocated_grid: Optional[aa.Grid2D] = None,
        mapper_list: Optional[aa.AbstractMapper] = None,
        operated_mapping_matrix: Optional[np.ndarray] = None,
        curvature_matrix_preload: Optional[np.ndarray] = None,
        curvature_matrix_counts: Optional[np.ndarray] = None,
        regularization_matrix: Optional[np.ndarray] = None,
        log_det_regularization_matrix_term: Optional[float] = None,
        failed=False,
    ):
        """
        Class which offers a concise API for settings up the preloads, which before a model-fit are set up via
        a comparison of two fits using two different models. If a quantity in these two fits is identical, it does
        not change thoughout the model-fit and can therefore be preloaded to avoid computation, speeding up
        the analysis.

        For example, the image-plane source-plane pixelization grid (which may be computationally expensive to compute
        via a KMeans algorithm) does not change for the majority of model-fits, because the associated model parameters
        are fixed. Preloading avoids rerruning the KMeans algorithm for every model fitted, by preloading it in memory
        and using this preload in every fit.

        Parameters
        ----------
        blurred_image
            The preloaded array of values containing the blurred image of a lens model fit (e.g. that light profile of
            every galaxy in the model). This can be preloaded when no light profiles in the model vary.
        w_tilde
            A class containing values that enable an inversion's linear algebra to use the w-tilde formalism. This can
            be preloaded when no component of the model changes the noise map (e.g. hyper galaxies are fixed).
        use_w_tilde
            Whether to use the w tilde formalism, which superseeds the value in `SettingsInversions` such that w tilde
            will be disabled for model-fits it is not applicable (e.g. because the noise-map changes).
        traced_grids_of_planes_for_inversion
            The two dimensional grids corresponding to the traced grids in a lens fit. This can be preloaded when no
             mass profiles in the model vary.
        sparse_image_plane_grid_list_of_planes
            The two dimensional grids corresponding to the sparse image plane grids in a lens fit, that is ray-traced to
            the source plane to form the source pixelization. This can be preloaded when no pixelizations in the model
            vary.
        relocated_grid
            The two dimensional grids corresponding to the grid that has had its border pixels relocated for a
            pixelization in a lens fit. This can be preloaded when no mass profiles in the model vary.
        mapper_list
            The mapper of a fit, which preloading avoids recalculation of the mapping matrix and image to source
            pixel mappings. This can be preloaded when no pixelizations in the model vary.
        operated_mapping_matrix
            A matrix containing the mappings between PSF blurred image pixels and source pixels used in the linear
            algebra of an inversion. This can be preloaded when no mass profiles and pixelizations in the model vary.
        curvature_matrix_preload
            A matrix containing preloaded value used to construct the curvature matrix from the blurred mapping matrix.
            This can be preloaded when no mass profiles and pixelizations in the model vary.
        curvature_matrix_counts
            A matrix containing the length of values in the curvature matrix preloaded, which are used to construct
            the curvature matrix from the blurred mapping matrix. This can be preloaded when no mass profiles and
            pixelizations in the model vary.

        Returns
        -------
        Preloads
            The preloads object used to skip certain calculations in the log likelihood function.
        """
        super().__init__(
            w_tilde=w_tilde,
            use_w_tilde=use_w_tilde,
            relocated_grid=relocated_grid,
            sparse_image_plane_grid_list_of_planes=sparse_image_plane_grid_list_of_planes,
            mapper_list=mapper_list,
            operated_mapping_matrix=operated_mapping_matrix,
            curvature_matrix_preload=curvature_matrix_preload,
            curvature_matrix_counts=curvature_matrix_counts,
            regularization_matrix=regularization_matrix,
            log_det_regularization_matrix_term=log_det_regularization_matrix_term,
        )

        self.blurred_image = blurred_image
        self.traced_grids_of_planes_for_inversion = traced_grids_of_planes_for_inversion
        self.failed = failed

    @classmethod
    def setup_all_via_fits(cls, fit_0, fit_1) -> "Preloads":
        """
        Setup the Preloads from two fits which use two different lens model of a model-fit.

        Parameters
        ----------
        fit_0
            The first fit corresponding to a model with a specific set of unit-values.
        fit_1
            The second fit corresponding to a model with a different set of unit-values.

        Returns
        -------
        Preloads
            Preloads which are set up based on the fit's passed in specific to a lens model.

        """

        preloads = cls()

        if isinstance(fit_0, aa.FitImaging):
            preloads.set_w_tilde_imaging(fit_0=fit_0, fit_1=fit_1)
            preloads.set_blurred_image(fit_0=fit_0, fit_1=fit_1)

        preloads.set_traced_grids_of_planes_for_inversion(fit_0=fit_0, fit_1=fit_1)
        preloads.set_sparse_image_plane_grid_list_of_planes(fit_0=fit_0, fit_1=fit_1)
        preloads.set_relocated_grid(fit_0=fit_0, fit_1=fit_1)
        preloads.set_mapper_list(fit_0=fit_0, fit_1=fit_1)
        preloads.set_operated_mapping_matrix_with_preloads(fit_0=fit_0, fit_1=fit_1)
        preloads.set_regularization_matrix_and_term(fit_0=fit_0, fit_1=fit_1)

        return preloads

    def set_blurred_image(self, fit_0, fit_1):
        """
        If the `LightProfile`'s in a model are all fixed parameters their corresponding image and therefore PSF blurred
        image do not change during the model fit and can therefore be preloaded.

        This function compares the blurred image of two fit's corresponding to two model instances, and preloads
        the blurred image if the blurred image of both fits are the same.

        The preload is typically used though out search chaining pipelines, as it is common to fix the lens light for
        the majority of model-fits.

        Parameters
        ----------
        fit_0
            The first fit corresponding to a model with a specific set of unit-values.
        fit_1
            The second fit corresponding to a model with a different set of unit-values.
        """
        self.blurred_image = None

        if np.max(abs(fit_0.blurred_image - fit_1.blurred_image)) < 1e-8:

            self.blurred_image = fit_0.blurred_image

            logger.info(
                "PRELOADS - Blurred image (e.g. the image of all light profiles) is preloaded for this model-fit."
            )

    def set_traced_grids_of_planes_for_inversion(self, fit_0, fit_1):
        """
        If the `MassProfiles`'s in a model are fixed their deflection angles and therefore corresponding traced grids
        do not change during the model-fit and can therefore be preloaded.

        This function compares the traced grids of two fit's corresponding to two model instances, and preloads the
        traced grids if the grids of both fits are the same. This preloaded grid is only used when constructing an
        inversion, because the `blurred_image` preload accounts for light profiles.

        The preload is typically used in hyper searches, where the mass model is fixed and the hyper-parameters are
        varied.

        Parameters
        ----------
        fit_0
            The first fit corresponding to a model with a specific set of unit-values.
        fit_1
            The second fit corresponding to a model with a different set of unit-values.
        """

        self.traced_grids_of_planes_for_inversion = None

        traced_grids_of_planes_0 = fit_0.tracer.traced_grids_of_planes_from(
            grid=fit_0.dataset.grid_inversion
        )

        traced_grids_of_planes_1 = fit_1.tracer.traced_grids_of_planes_from(
            grid=fit_1.dataset.grid_inversion
        )

        if traced_grids_of_planes_0[-1] is not None:

            if (
                traced_grids_of_planes_0[-1].shape[0]
                == traced_grids_of_planes_1[-1].shape[0]
            ):

                if (
                    np.max(
                        abs(traced_grids_of_planes_0[-1] - traced_grids_of_planes_1[-1])
                    )
                    < 1e-8
                ):

                    self.traced_grids_of_planes_for_inversion = traced_grids_of_planes_0

                    logger.info(
                        "PRELOADS - Traced grid of planes (for inversion) preloaded for this model-fit."
                    )

    def set_sparse_image_plane_grid_list_of_planes(self, fit_0, fit_1):
        """
        If the `Pixelization`'s in a model are fixed their image-plane sparse grid (which defines the set of pixels
        that are ray-traced to construct the source-plane pixelization) do not change during the model=fit and
        can therefore be preloaded.

        This function compares the image plane sparse grid of two fit's corresponding to two model instances, and p
        reloads the grid if the grids of both fits are the same.

        The preload is typically used thoughout search chaining pipelines which use inversions, as it is common to
        for the pixelization's parameters to only vary in the hyper-searches.

        Parameters
        ----------
        fit_0
            The first fit corresponding to a model with a specific set of unit-values.
        fit_1
            The second fit corresponding to a model with a different set of unit-values.
        """

        self.sparse_image_plane_grid_list_of_planes = None

        sparse_image_plane_grid_list_of_planes_0 = fit_0.tracer.sparse_image_plane_grid_list_of_planes_from(
            grid=fit_0.dataset.grid_inversion
        )

        sparse_image_plane_grid_list_of_planes_1 = fit_1.tracer.sparse_image_plane_grid_list_of_planes_from(
            grid=fit_1.dataset.grid_inversion
        )

        if sparse_image_plane_grid_list_of_planes_0[-1] is not None:

            if sparse_image_plane_grid_list_of_planes_0[-1][0] is not None:

                if (
                    sparse_image_plane_grid_list_of_planes_0[-1][0].shape[0]
                    == sparse_image_plane_grid_list_of_planes_1[-1][0].shape[0]
                ):

                    if (
                        np.max(
                            abs(
                                sparse_image_plane_grid_list_of_planes_0[-1][0]
                                - sparse_image_plane_grid_list_of_planes_1[-1][0]
                            )
                        )
                        < 1e-8
                    ):

                        self.sparse_image_plane_grid_list_of_planes = (
                            sparse_image_plane_grid_list_of_planes_0
                        )

                        logger.info(
                            "PRELOADS - Sparse image-plane grids of planes is preloaded for this model-fit."
                        )

    def output_info_to_summary(self, file_path):

        file_preloads = path.join(file_path, "preloads.summary")

        if self.failed:

            logger.info(
                "PRELOADS - Preloading failed as models gave too many exceptions, preloading therefore "
                "not used."
            )

            af.formatter.output_list_of_strings_to_file(
                file=file_preloads, list_of_strings=["FAILED"]
            )

        else:

            af.formatter.output_list_of_strings_to_file(
                file=file_preloads, list_of_strings=self.info
            )

    def reset_all(self):
        """
        Reset all preloads, typically done at the end of a model-fit to save memory.
        """
        self.w_tilde = None

        self.blurred_image = None
        self.traced_grids_of_planes_for_inversion = None
        self.sparse_image_plane_grid_list_of_planes = None
        self.relocated_grid = None
        self.mapper = None
        self.blurred_mapping_matrix = None
        self.curvature_matrix_preload = None
        self.curvature_matrix_counts = None
        self.regularization_matrix = None
        self.log_det_regularization_matrix_term = None

    @property
    def info(self) -> List[str]:
        """
        The information on what has or has not been preloaded, which is written to the file `preloads.summary`.

        Returns
        -------
            A list of strings containing True and False values as to whether a quantity has been preloaded.
        """
        line = [f"W Tilde = {self.w_tilde is not None}\n"]
        line += [f"Use W Tilde = {self.use_w_tilde}\n\n"]
        line += [f"Blurred Image = {np.count_nonzero(self.blurred_image) != 0}\n"]
        line += [
            f"Traced Grids of Planes (For LinearEqn) = {self.traced_grids_of_planes_for_inversion is not None}\n"
        ]
        line += [
            f"Sparse Image-Plane Grids of Planes = {self.sparse_image_plane_grid_list_of_planes is not None}\n"
        ]
        line += [f"Relocated Grid = {self.relocated_grid is not None}\n"]
        line += [f"Mapper = {self.mapper_list is not None}\n"]
        line += [
            f"Blurred Mapping Matrix = {self.operated_mapping_matrix is not None}\n"
        ]
        line += [
            f"Curvature Matrix Sparse = {self.curvature_matrix_preload is not None}\n"
        ]
        line += [f"Regularization Matrix = {self.regularization_matrix is not None}\n"]
        line += [
            f"Log Det Regularization Matrix Term = {self.log_det_regularization_matrix_term is not None}\n"
        ]

        return line


"""Serialization utilities for Alpa.
Adapted from https://flax.readthedocs.io/en/latest/_modules/flax/serialization.html.
Add support for DistributedArray and ReplicatedDistributedArray serialization in Alpa.
"""

import enum
import logging
import os
import re
from typing import Union, Any, Sequence
import uuid

import numpy as np
import jax
from jax.interpreters.pxla import ShardingSpec
from jax.core import ShapedArray
import jax.numpy as jnp
from jax._src.tree_util import tree_flatten, tree_leaves, tree_unflatten
from flax.serialization import to_state_dict, from_state_dict, _ndarray_from_bytes, _ndarray_to_bytes
import msgpack
import tensorstore as ts

from alpa.device_mesh import DistributedArray, ReplicatedDistributedArray, PhysicalDeviceMesh

PyTree = Any

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class _MsgpackExtType(enum.IntEnum):
    """Messagepack custom type ids."""
    ndarray = 1
    native_complex = 2
    npscalar = 3
    distarray = 4
    replicated_distarray = 5

def _msgpack_ext_pack_wrapper(ckpt_dir):

    def _get_save_path():
        return os.path.join(ckpt_dir, uuid.uuid4().hex)

    def _msgpack_ext_pack(x):
        """Messagepack encoders for custom types."""
        if isinstance(x, (np.ndarray, jax.xla.DeviceArray)):
            save_dir = _get_save_path()
            ts_store(save_dir, x)
            return msgpack.ExtType(_MsgpackExtType.ndarray,
                                   msgpack.packb(save_dir))
        if np.issctype(type(x)):
            # pack scalar as ndarray
            return msgpack.ExtType(_MsgpackExtType.npscalar,
                                   _ndarray_to_bytes(np.asarray(x)))
        elif isinstance(x, complex):
            return msgpack.ExtType(_MsgpackExtType.native_complex,
                                   msgpack.packb((x.real, x.imag)))
        elif isinstance(x, DistributedArray):
            save_dir = _get_save_path()
            x.save(save_dir)
            return msgpack.ExtType(_MsgpackExtType.distarray,
                                   msgpack.packb(save_dir))
        elif isinstance(x, ReplicatedDistributedArray):
            save_dir = _get_save_path()
            x.replica.save(save_dir)
            return msgpack.ExtType(_MsgpackExtType.replicated_distarray,
                                   msgpack.packb(save_dir))
        return x

    return _msgpack_ext_pack

def _msgpack_ext_unpack(code, data):
    """Messagepack decoders for custom types."""
    if code == _MsgpackExtType.ndarray:
        return msgpack.unpackb(data)
    elif code == _MsgpackExtType.native_complex:
        complex_tuple = msgpack.unpackb(data)
        return complex(complex_tuple[0], complex_tuple[1])
    elif code == _MsgpackExtType.npscalar:
        ar = _ndarray_from_bytes(data)
        return ar[()]  # unpack ndarray to scalar
    elif code == _MsgpackExtType.distarray:
        return msgpack.unpackb(data)
    elif code == _MsgpackExtType.replicated_distarray:
        return msgpack.unpackb(data)
    return msgpack.ExtType(code, data)

def get_ts_spec(ckpt_path: str):
    spec = {
        'driver': 'zarr',
        'kvstore': {},
        'metadata_key': ".zarray0"
    }
    if ckpt_path.startswith('gs://'):
        m = re.fullmatch('^gs://([^/]*)/(.*)$', ckpt_path, re.DOTALL)
        if m is None:
            raise ValueError(
                'The ckpt_path should contain the bucket name and the '
                f'file path inside the bucket. Got: {ckpt_path}')
        gcs_bucket = m.group(1)
        path_without_bucket = m.group(2)
        spec['kvstore'] = {
            'driver': 'gcs',
            'bucket': gcs_bucket,
            'path': path_without_bucket
        }
    else:
        spec['kvstore'] = {'driver': 'file', 'path': ckpt_path}
    return spec

def ts_store(ckpt_dir, data: Union[np.ndarray, jax.xla.DeviceArray]):
    ts_spec = get_ts_spec(ckpt_dir)
    dtype = data.dtype
    if dtype == jnp.bfloat16:
        # Tensorstore uses 'bfloat16', not '<V2'.
        dtype = 'bfloat16'
    else:
        dtype = np.dtype(dtype).str
    metadata = {
        'shape': data.shape,
        'chunks': data.shape,
        'dtype': dtype,
    }
    ts_spec['metadata'] = metadata
    t = ts.open(ts.Spec(ts_spec),
                create=True,
                open=True,
                context=ts.Context({'file_io_concurrency': {
                    'limit': 128
                }})).result()

    t.write(data).result()

def save_checkpoint(ckpt_dir: Union[str, os.PathLike], target: PyTree,
                    step: int):
    """Save a checkpoint of the `target` to `path`.

        Similar to flax.training.checkpoints.save_checkpoint, but support DistributedArrays
        and ReplicatedDistributedArray in alpa. Also it will save np.ndarray and jax.xla.DeviceArray
        into tensorstore for later distributed loading.
        # TODO (zhongyinmin): copy all the safe-saving stuff from
        https://flax.readthedocs.io/en/latest/_modules/flax/training/checkpoints.html#save_checkpoint

        Args:
           ckpt_dir: str or pathlib-like path to store checkpoint directories in.
           target: serializable flax object, usually a trainState
           step: training step number or other metric number
    """
    state_dict = to_state_dict(target)
    os.makedirs(ckpt_dir, exist_ok=True)
    ckpt_path = os.path.join(ckpt_dir, f"checkpoint_{step}")
    with open(ckpt_path, 'wb') as fp:
        fp.write(
            msgpack.packb(state_dict,
                          default=_msgpack_ext_pack_wrapper(ckpt_dir),
                          strict_types=True))

class LoadInfo:
    """
    A wrapper for the loading information.
    """
    def __init__(self,
                 avals: Sequence[ShapedArray],
                 meshes: Sequence[PhysicalDeviceMesh],
                 specs: Sequence[ShardingSpec]):
        assert len(avals) == len(meshes)
        assert len(meshes) == len(specs)
        self.avals = avals
        self.meshes = meshes
        self.specs = specs

    def add_replica(self, aval, mesh, spec):
        self.avals.append(aval)
        self.meshes.append(mesh)
        self.specs.append(spec)

    def get_info(self):
        if self.is_replicated():
            return zip(self.avals, self.meshes, self.specs)
        else:
            return self.avals[0], self.meshes[0], self.specs[0]

    def is_replicated(self):
        return len(self.avals) > 1

    def __str__(self):
        return f"{self.avals[0]}, {self.meshes[0].mesh_id}, {self.specs[0]}"

def restore_checkpoint(ckpt_dir: Union[str, os.PathLike], step: int, load_info: PyTree):
    """Restore the specified checkpoint from `path`.

        Similar to flax.training.checkpoints.load_checkpoint,
        but support DistributedArrays and ReplicatedDistributedArray in alpa.
        # TODO (zhongyinmin): copy all the safe-loading stuff from
        https://flax.readthedocs.io/en/latest/_modules/flax/training/checkpoints.html#restore_checkpoint

        Args:
            ckpt_dir: directory of checkpoints to restore from.
            step: step number to load.
            load_info: shardingSpec and deviceMesh allocation info for loading.
    """
    ckpt_path = os.path.join(ckpt_dir, f"checkpoint_{step}")
    with open(ckpt_path, 'rb') as fp:
        ckpt_contents = fp.read()
    state_dict_content = msgpack.unpackb(ckpt_contents,
                                         ext_hook=_msgpack_ext_unpack,
                                         raw=False)
    state_paths, state_tree = tree_flatten(from_state_dict(load_info, state_dict_content))
    flat_info = tree_leaves(load_info)
    flat_load_state = []
    for path, info in zip(state_paths, flat_info):
        if info is None:
            logger.warning("Variable is not used, skip loading it")
            flat_load_state.append(None)
        if info.is_replicated():
            meshes, arrays = [], []
            for aval, mesh, spec in info.get_info():
                meshes.append(mesh)
                arrays.append(DistributedArray.load(path, aval, mesh, spec))
            flat_load_state.append(ReplicatedDistributedArray(meshes, arrays))
        else:
            aval, mesh, spec = info.get_info()
            flat_load_state.append(DistributedArray.load(path, aval, mesh, spec))
    return tree_unflatten(state_tree, flat_load_state)


#!/usr/bin/env python

import pyaudio
import numpy
import wave
import time
import sys
from pprint import pprint

tunings = {
        'standard': [
            ('E', 329.63),
            ('B', 246.94),
            ('G', 196.00),
            ('D', 146.83),
            ('A', 110.00),
            ('E', 82.41)
        ]
}

AUDIO_CHANNELS = 1
AUDIO_SAMP_FREQ = 48000
AUDIO_FRAMES_PER_BUFFER = 1024
AUDIO_SAMP_DELAY = 0.5

def freq(frames, sampling_frequency=44100, frames_per_buffer=1024):
    # unpack wave and multiply by the window
    frame_array = numpy.array(wave.struct.unpack("%dh" % (len(frames) / 2), frames)) * numpy.blackman(frames_per_buffer)

    # calculate fft and maximum
    fftout = abs(numpy.fft.rfft(frame_array)) ** 2
    maximum = fftout[1:].argmax() + 1

    if maximum != len(fftout) - 1:
        y0, y1, y2 = numpy.log(fftout[(maximum - 1):(maximum + 2):])
        x1 = (y2 - y0) * 0.5 / (2 * y1 - y2 - y0)
        result = (maximum + x1) * sampling_frequency / frames_per_buffer
    else:
        result = maximum * sampling_frequency / frames_per_buffer

    return result

def main():
    audio = pyaudio.PyAudio()
    stream = audio.open(format=pyaudio.paInt16,
            channels=AUDIO_CHANNELS,
            rate=AUDIO_SAMP_FREQ,
            input=True,
            frames_per_buffer=AUDIO_FRAMES_PER_BUFFER)

    pprint(tunings['standard'])
    print
    while(True):
        buff = stream.read(AUDIO_FRAMES_PER_BUFFER)
        sys.stdout.write( '\r' + str(freq(buff, AUDIO_SAMP_FREQ, AUDIO_FRAMES_PER_BUFFER)) )
        sys.stdout.flush()
        time.sleep(AUDIO_SAMP_DELAY)

    stream.stop_stream()
    stream.close()
    audio.terminate()

if __name__ == '__main__':
    main()


"""
Created on Feb 24, 2018

@author: Siyuan Qi

Description of the file.

"""

import os
import time
import random
import pickle
import argparse
import warnings

import cv2
import torch.utils.data
import numpy as np
import vsrl_utils as vu

from . import vcoco_config
from . import metadata
from . import roi_pooling

import torchvision
import torchvision.transforms

import matplotlib.pyplot as plt
from skimage import io, transform, color

class Vgg16(torch.nn.Module):
    def __init__(self, last_layer=0, requires_grad=False):
        super(Vgg16, self).__init__()
        pretrained_vgg = torchvision.models.vgg16(pretrained=True)
        self.features = torch.nn.Sequential()
        for x in range(len(pretrained_vgg.features)):
            self.features.add_module(str(x), pretrained_vgg.features[x])

        # if feature_mode == 'roi_vgg':
        #     self.classifier = torch.nn.Sequential()
        #     self.classifier.add_module(str(0), pretrained_vgg.classifier[0])

        # self.classifier.add_module(str(1), pretrained_vgg.classifier[1])
        # for x in range(len(pretrained_vgg.classifier)-last_layer):
        #     print pretrained_vgg.classifier[x]
        #     self.classifier.add_module(str(x), pretrained_vgg.classifier[x])

        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, x):
        x = self.features(x)

        # if feature_mode == 'roi_vgg':
        #     x = x.view(x.size(0), -1)
        #     x = self.classifier(x)
        return x

def get_model():
    feature_network = Vgg16(last_layer=1)
    feature_network.cuda()

    transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225]),
    ])

    return feature_network, transform

def combine_box(box1, box2):
    return np.hstack((np.maximum(box1[:2], box2[:2]), np.minimum(box1[2:], box2[2:])))

mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])

def img_to_torch(img, t):
    """
    input: H x W x C img iterables with range 0-255
    output: C x H x W img tensor with range 0-1, normalized
    """
    img = np.array(img) / 255.
    img = (img - mean) / std
    if len(img.shape) == 3:
        img = np.expand_dims(img.transpose([2,0,1]), axis=0)
    elif len(img.shape) == 4:
        img = img.transpose([0,3,1,2])
    elif len(img.shape) == 5:
        img = img.transpose([0,1,4,2,3])
    img = torch.autograd.Variable(torch.Tensor(img)).cuda()
    return img

def transform_bbox(boxes, mat):
    bboxes2 = np.zeros([boxes.shape[0], 4, 2])
    bboxes2[:,:2,0] = boxes[:,[2]]
    bboxes2[:,2:,0] = boxes[:,[0]]
    bboxes2[:,::2,1] = boxes[:,[3]]
    bboxes2[:,1::2,1] = boxes[:,[1]]
    bboxes_transformed = np.zeros(bboxes2.shape)
    bboxes_transformed[:,:,0] = bboxes2[:,:,0] * mat[0,0] + bboxes2[:,:,1] * mat[0,1] + mat[0,2] # X = a0*x + a1*y + a2
    bboxes_transformed[:,:,1] = bboxes2[:,:,0] * mat[1,0] + bboxes2[:,:,1] * mat[1,1] + mat[1,2] # Y = b0*x + b1*y + b2
    boxes = np.zeros([boxes.shape[0], 4])
    boxes[:,2] = np.min(bboxes_transformed[:,:,0], axis=1)
    boxes[:,0] = np.max(bboxes_transformed[:,:,0], axis=1)
    boxes[:,3] = np.min(bboxes_transformed[:,:,1], axis=1)
    boxes[:,1] = np.max(bboxes_transformed[:,:,1], axis=1)
    return boxes

class NoisyVCOCO(torch.utils.data.Dataset):
    def __init__(self, root, imageset, node_feature_appd=False, chance=0.0):
        self.root = root
        self.coco = vu.load_coco(root)
        vcoco_all = vu.load_vcoco('vcoco_{}'.format(imageset), root)
        self.image_ids = vcoco_all[0]['image_id'][:, 0].astype(int).tolist()
        self.unique_image_ids = list(set(self.image_ids))
        self.node_feature_appd = node_feature_appd
        self.chance = chance

        self.model, self.transform = get_model()

    def __getitem__(self, index):
        input_h, input_w = 224, 224
        feature_length = 4096
        feature_size = [7, 7]
        adaptive_max_pool = roi_pooling.AdaptiveMaxPool2d(*feature_size)

        img_name = self.coco.loadImgs(ids=[self.unique_image_ids[index]])[0]['file_name']
        img_type = img_name.split('_')[1]
        try:
            # img_path = os.path.join(self.root, '../coco/coco/images', '{}2014'.format(img_type), img_name)
            # data = pickle.load(open(os.path.join(self.root, '..', 'processed', 'resnet', '{}.p'.format(img_name)), 'rb'))
            img_path = os.path.join('/media/tengyu/data/mscoco/%s/%s'%(img_type, img_name))
            data = pickle.load(open(os.path.join(os.path.dirname(__file__), '../../../../data/feature_resnet_noisy', '{}.p'.format(img_name)), 'rb'))
            img = io.imread(img_path)
            if len(img.shape) == 2:
                img = np.tile(np.expand_dims(img, axis=-1), [1,1,3])
        except IOError:
            warnings.warn('data missing for {}'.format(img_name))
            return self.__getitem__(3)
        except ValueError:
            warnings.warn('img grayscale for {}'.format(img_name))
            return self.__getitem__(3)

        h, w, _ = img.shape
        img_id = data['img_id']
        adj_mat = data['adj_mat']
        node_labels = data['node_labels']
        node_roles = data['node_roles']
        obj_boxes = data['obj_boxes']
        part_boxes = data['part_boxes']
        human_boxes = data['human_boxes']
        human_num = data['human_num']
        part_num = data['part_num']
        obj_num = data['obj_num']
        obj_classes = data['obj_classes']
        part_classes = data['part_classes']
        part_human_id = data['part_human_id']
        part_adj_mat = None # data['part_adj_mat']

        part_images = []
        obj_images = []
        edge_images = []

        # apply transformation
        if random.random() < self.chance:
            shear = random.random() * 30 - 15        # -15 deg to +15 deg
            shear = np.deg2rad(shear)
        else:
            shear = 0
        if random.random() < self.chance:
            rotate = random.random() * 30 - 15       # -15 deg to +15 deg
            rotate = np.deg2rad(rotate)
        else:
            rotate = 0
        if random.random() < self.chance:
            scale = [random.random() * 0.3 + 0.85, random.random() * 0.3 + 0.85]     # 0.85 to 1.15
        else:
            scale = [1.0, 1.0]

        translation = [img.shape[0] * 0.5, img.shape[1] * 0.5]

        if random.random() < self.chance:
            # flip
            img = img[:,::-1,:]
            for b in part_boxes:
                b[0] = w - b[0]
                b[2] = w - b[2]
            for b in obj_boxes:
                b[0] = w - b[0]
                b[2] = w - b[2]

        img = color.rgb2hsv(img)
        if random.random() < self.chance:
            # brightness
            img[:,:,2] += random.random() * 0.3 - 0.15  # -0.15 to + 0.15
            img = np.clip(img, 0, 1)
        if random.random() < self.chance:
            # hue
            img[:,:,0] += random.random() * 0.5 - 0.25  # -0.15 to + 0.15
            img = np.clip(img, 0, 1)
        img = color.hsv2rgb(img)

        transformation = transform.AffineTransform(shear=shear, scale=scale, rotation=rotate, translation=translation)
        mat = transformation.params

        img = (img * 255).astype(np.uint8)
        img = transform.warp(img, transformation.inverse, order=0, output_shape=(img.shape[0] * 2, img.shape[1] * 2))

        part_boxes = transform_bbox(part_boxes, mat)
        obj_boxes = transform_bbox(obj_boxes, mat)

        img_feature = self.model.features(img_to_torch(cv2.resize(img, (input_h, input_w), interpolation=cv2.INTER_LINEAR), self.transform))
        img_feature = torch.autograd.Variable(img_feature)
        scale_h = (feature_size[0] - 1)/float(img.shape[0])
        scale_w = (feature_size[1] - 1)/float(img.shape[1])

        obj_features = np.empty((obj_num, 512 * feature_size[0] * feature_size[1]))
        objs = np.copy(obj_boxes)
        objs[:, 0] *= scale_w
        objs[:, 2] *= scale_w
        objs[:, 1] *= scale_h
        objs[:, 3] *= scale_h

        for i_box in range(objs.shape[0]):
            obj = objs[i_box, :].astype(int)
            if (obj[1] - obj[3]) < 1 or (obj[0] - obj[2]) < 1:
                continue
            try:
                obj_feature = adaptive_max_pool(img_feature[..., obj[3]:(obj[1]), obj[2]:(obj[0])])
                obj_features[i_box, :] = obj_feature.data.cpu().numpy().flatten()
            except:
                print(obj)
                raise

        part_features = np.empty((part_num, 512 * feature_size[0] * feature_size[1]))
        parts = np.copy(part_boxes)
        parts[:, 0] *= scale_w
        parts[:, 2] *= scale_w
        parts[:, 1] *= scale_h
        parts[:, 3] *= scale_h

        for i_box in range(parts.shape[0]):
            part = parts[i_box, :].astype(int)
            if (part[1] - part[3]) < 1 or (part[0] - part[2]) < 1:
                continue
            try:
                part_feature = adaptive_max_pool(img_feature[..., part[3]:(part[1]), part[2]:(part[0])])
                part_features[i_box, :] = part_feature.data.cpu().numpy().flatten()
            except:
                print(part)
                raise

        edge_features = np.empty((part_num + obj_num, part_num + obj_num, 512 * feature_size[0] * feature_size[1]))

        for i_part, part_box in enumerate(part_boxes):
            edge_images.append([])
            for i_obj, obj_box in enumerate(obj_boxes):
                edge_box = combine_box(obj_box, part_box).astype(int)
                edge_box[0] *= scale_w
                edge_box[2] *= scale_w
                edge_box[1] *= scale_h
                edge_box[3] *= scale_h

                edge_box = edge_box.astype(int)
                if (edge_box[1] - edge_box[3]) < 1 or (edge_box[0] - edge_box[2]) < 1:
                    continue
                try:
                    edge_feature = adaptive_max_pool(img_feature[..., edge_box[3]:(edge_box[1]), edge_box[2]:(edge_box[0])])
                    edge_features[i_part, part_num + i_obj, :] = edge_feature.data.cpu().numpy().flatten()
                    edge_features[part_num + i_obj, i_part, :] = edge_feature.data.cpu().numpy().flatten()
                except:
                    print(edge_box)
                    raise

        node_features = np.concatenate([part_features, obj_features], axis=0)

        # append bbox and class to node features
        if self.node_feature_appd:
            part_eye = np.eye(14)
            obj_eye = np.eye(81)

            node_features_appd = np.zeros([node_features.shape[0], 6 + 14 + 81])

            node_features_appd[:part_num,0] = (part_boxes[:,2] - part_boxes[:,0]) / img_w # relative w
            node_features_appd[:part_num,1] = (part_boxes[:,3] - part_boxes[:,1]) / img_h # relative h
            node_features_appd[:part_num,2] = ((part_boxes[:,2] + part_boxes[:,0]) / 2) / img_w # relative cx
            node_features_appd[:part_num,3] = ((part_boxes[:,3] + part_boxes[:,1]) / 2) / img_h # relative cy
            node_features_appd[:part_num,4] = (part_boxes[:,2] - part_boxes[:,0]) * (part_boxes[:,3] - part_boxes[:,1]) / (img_w * img_h) # relative area
            node_features_appd[:part_num,5] = (part_boxes[:,2] - part_boxes[:,0]) / (part_boxes[:,3] - part_boxes[:,1]) # aspect ratio
            node_features_appd[:part_num,6:6+14] = part_eye[part_classes]

            node_features_appd[part_num:,0] = (obj_boxes[:,2] - obj_boxes[:,0]) / img_w # relative w
            node_features_appd[part_num:,1] = (obj_boxes[:,3] - obj_boxes[:,1]) / img_h # relative h
            node_features_appd[part_num:,2] = ((obj_boxes[:,2] + obj_boxes[:,0]) / 2) / img_w # relative cx
            node_features_appd[part_num:,3] = ((obj_boxes[:,3] + obj_boxes[:,1]) / 2) / img_h # relative cy
            node_features_appd[part_num:,4] = (obj_boxes[:,2] - obj_boxes[:,0]) * (obj_boxes[:,3] - obj_boxes[:,1]) / (img_w * img_h) # relative area
            node_features_appd[part_num:,5] = (obj_boxes[:,2] - obj_boxes[:,0]) / (obj_boxes[:,3] - obj_boxes[:,1]) # aspect ratio
            node_features_appd[part_num:,6+14:] = obj_eye[obj_classes]

            node_features_appd[np.isnan(node_features_appd)] = 0
            node_features_appd[np.isinf(node_features_appd)] = 0

            node_features = np.concatenate([node_features, node_features_appd], axis=-1)

            node_features[np.isnan(node_features)] = 0
            node_features[np.isinf(node_features)] = 0
            edge_features[np.isnan(edge_features)] = 0
            edge_features[np.isinf(edge_features)] = 0

        return edge_features, \
                    node_features, \
                    part_human_id, \
                    adj_mat, \
                    node_labels, \
                    node_roles, \
                    obj_boxes, \
                    part_boxes, \
                    human_boxes, \
                    img_id, \
                    img_name, \
                    human_num, \
                    part_num, \
                    obj_num, \
                    obj_classes, \
                    part_classes, \
                    part_adj_mat, \
                    img_name

    def __len__(self):
        return len(self.unique_image_ids)

def main(args):
    subset = ['train', 'val', 'test']
    training_set = NoisyVCOCO(args.data_root, subset[0])
    print('{} instances.'.format(len(training_set)))
    start_time = time.time()
    edge_features, obj_features, part_features, part_human_id, adj_mat, obj_labels, obj_roles, human_labels, human_roles, obj_boxes, part_boxes, human_boxes, img_id, img_name, human_num, obj_num, obj_classes, part_classes = training_set[113]
    print('Time elapsed: {:.2f}s'.format(time.time() - start_time))

    start_time = time.time()
    edge_features, obj_features, part_features, part_human_id, adj_mat, obj_labels, obj_roles, human_labels, human_roles, obj_boxes, part_boxes, human_boxes, img_id, img_name, human_num, obj_num, obj_classes, part_classes = training_set[113]
    print('Time elapsed: {:.2f}s'.format(time.time() - start_time))

    start_time = time.time()
    edge_features, obj_features, part_features, part_human_id, adj_mat, obj_labels, obj_roles, human_labels, human_roles, obj_boxes, part_boxes, human_boxes, img_id, img_name, human_num, obj_num, obj_classes, part_classes = training_set[113]
    print('Time elapsed: {:.2f}s'.format(time.time() - start_time))

    start_time = time.time()
    edge_features, obj_features, part_features, part_human_id, adj_mat, obj_labels, obj_roles, human_labels, human_roles, obj_boxes, part_boxes, human_boxes, img_id, img_name, human_num, obj_num, obj_classes, part_classes = training_set[113]
    print('Time elapsed: {:.2f}s'.format(time.time() - start_time))

    start_time = time.time()
    edge_features, obj_features, part_features, part_human_id, adj_mat, obj_labels, obj_roles, human_labels, human_roles, obj_boxes, part_boxes, human_boxes, img_id, img_name, human_num, obj_num, obj_classes, part_classes = training_set[113]
    print('Time elapsed: {:.2f}s'.format(time.time() - start_time))

def parse_arguments():
    paths = vcoco_config.Paths()
    parser = argparse.ArgumentParser(description='V-COCO dataset')
    parser.add_argument('--data-root', default=paths.data_root, help='dataset path')
    parser.add_argument('--tmp-root', default=paths.tmp_root, help='intermediate result path')
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_arguments()
    main(args)


import torch
import numpy as np
import pandas as pd
import librosa as lr
import matplotlib.pyplot as plt

from librosa import display as lrd
import IPython.display as ipd

from torch.utils.data import DataLoader, ConcatDataset, random_split
from asteroid.data import TimitDataset
from tqdm import tqdm

from asteroid.data.utils import find_audio_files, cut_or_pad

from torch import optim
from pytorch_lightning import Trainer, loggers as pl_loggers
from asteroid.engine import System
from asteroid.losses import singlesrc_neg_sisdr

from asteroid.engine.system import UNetGAN

def sisdr_loss_wrapper(est_target, target):
    return singlesrc_neg_sisdr(est_target.squeeze(1), target).mean()

def train_val_split(ds, val_fraction=0.1, random_seed=42):
    assert val_fraction > 0 and val_fraction < 0.5
    len_train = int(len(ds) * (1 - val_fraction))
    len_val = len(ds) - len_train
    return random_split(ds, [len_train, len_val], generator=torch.Generator().manual_seed(random_seed))

def main():
    TIMIT_CACHE_DIR = '/import/vision-eddydata/dm005_tmp/mixed_wavs_asteroid'
    train_snrs = [-25, -20, -15, -10, -5, 0, 5, 10, 15]
    test_snrs = [-30, -25, -20, -15, -10, -5, 0]

    timit_train_misc = TimitDataset.load_with_cache(
        '../../../datasets/TIMIT', '../../../datasets/noises-train',
        cache_dir=TIMIT_CACHE_DIR, snrs=train_snrs, root_seed=42, prefetch_mixtures=False,
        dset_name='train-misc', subset='train', track_duration=48000)

    train_set, val_set = train_val_split(timit_train_misc)

    BATCH_SIZE = 64
    NUM_WORKERS = 10

    train_loader = DataLoader(
        train_set,
        shuffle=True,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        drop_last=True,
    )

    val_loader = DataLoader(
        val_set,
        shuffle=False,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        drop_last=True,
    )

    unetgan_module = UNetGAN(lr_g=2e-4, lr_d=2e-4)

    trainer = Trainer(max_epochs=102, gpus=-1, accelerator='dp',
                      resume_from_checkpoint='logs/lightning_logs/version_0/checkpoints/epoch=99-step=395190.ckpt')
    trainer.fit(unetgan_module, train_loader, val_loader)

if __name__ == "__main__":
    main()


# Copyright 2020 Pulser Development Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The Sequence class, where a pulse sequence is defined."""

from __future__ import annotations

from collections import namedtuple
from collections.abc import Callable, Generator, Iterable, Set
import copy
from functools import wraps
from itertools import chain
import json
from sys import version_info
from typing import Any, Dict, cast, NamedTuple, Optional, Tuple, Union
# from typing_extensions import ParamSpecKwargs
import warnings

import numpy as np
from numpy.ma import diff
from numpy.typing import ArrayLike
from qutip.states import projection

from braket.analog.quera.lattice import Lattice
from braket.analog.quera.channels import Channel, Hyperfine, Rydberg
from braket.analog.quera.waveforms import Waveform, ConstantWaveform, CustomWaveform
from braket.analog.quera._seq_drawer import draw_sequence

# import pulser
# from pulser.channels import Channel
# from pulser.devices import MockDevice
# from pulser.devices._device_datacls import Device
# from pulser.json.coders import PulserEncoder, PulserDecoder
# from pulser.json.utils import obj_to_dict
# from pulser.parametrized import Parametrized, Variable
# from pulser.pulse import Pulse
# from pulser.register import Register

# if version_info[:2] >= (3, 8):  # pragma: no cover
#     from typing import Literal, get_args
# else:  # pragma: no cover
#     try:
#         from typing_extensions import Literal, get_args  # type: ignore
#     except ImportError:
#         raise ImportError(
#             "Using pulser with Python version 3.7 requires the"
#             " `typing_extensions` module. Install it by running"
#             " `pip install typing-extensions`."
#         )

QubitId = Union[int, str]

class _TimeSlot(NamedTuple):
    """Auxiliary class to store the information in the schedule."""

    type: Union[Tuple, str] # Union[Pulse, str]
    # Ω: Waveform
    # Δ: Waveform
    # ϕ: Waveform
    ti: int
    tf: int
    targets: Set[QubitId]

# Encodes a sequence building calls
_Call = namedtuple("_Call", ["name", "args", "kwargs"])

class Sequence:
    def __init__(self, lattice: Lattice):
        """Initializes a new pulse sequence."""
        if not isinstance(lattice, Lattice):
            raise TypeError("'lattice' must be of type 'Lattice'.")

        self._lattice: Lattice = lattice
        self._channels: dict[str, Channel] = {}
        self._qids: set[QubitId] = set(self.qubit_info.keys())
        # self._duration: int = 0
        self._schedule: dict[str, list[_TimeSlot]] = {}
        # self._calls: list[_Call] = [_Call("__init__", (lattice), {})]
        # Last time a target happened, by channel
        self._last_target: dict[str, int] = {}

    @property
    def qubit_info(self) -> dict[QubitId, np.ndarray]:
        """Dictionary with the qubit's IDs and positions."""
        return self._lattice.qubits

    @property
    def declared_channels(self) -> dict[str, Channel]:
        """Channels declared in this Sequence."""
        return dict(self._channels)

    def declare_channel(
        self,
        channel_name: str,
        channel_id: str,
        initial_target: Optional[Union[Iterable[QubitId],QubitId]] = None,
    ) -> None:
        """Declares a new channel to the Sequence.

        The first declared channel implicitly defines the sequence's mode of
        operation (i.e. the underlying Hamiltonian). In particular, if the
        first declared channel is of type ``Microwave``, the sequence will work
        in "XY Mode" and will not allow declaration of channels that do not
        address the 'XY' basis. Inversely, declaration of a channel of another
        type will block the declaration of ``Microwave`` channels.

        Note:
            Regular devices only allow a channel to be declared once, but
            ``MockDevice`` channels can be repeatedly declared if needed.

        Args:
            channel_name (str): Unique channel_name for the channel in the sequence.
            channel_id (str): How the channel is identified in the device.
                Consult ``Sequence.available_channels`` to see which channel
                ID's are still available and the associated channel's
                description.
            initial_target (Optional[Union[int, str, Iterable]]): For 'Local'
                addressing channels only. Declares the initial target of the
                channel. If left as None, the initial target will have to be
                set manually as the first addition to this channel.
        """
        if channel_name in self._channels:
            raise ValueError("The given channel name is already in use.")

        # if channel_id not in self._device.channels:
        if channel_id not in ['Hyperfine', 'Rydberg']:
            raise ValueError("No channel %s in the device." % channel_id)

        ch = Rydberg() if channel_id == "Rydberg" else Hyperfine()

        self._channels[channel_name] = ch
        self._schedule[channel_name] = []
        self._last_target[channel_name] = 0

        if initial_target == None:
            self._add_to_schedule(channel_name, _TimeSlot("target", -1, 0, self._qids))
        else:
            self._target(
                cast(Union[Iterable, QubitId], initial_target), channel_name
            )

        # # Manually store the channel declaration as a regular call
        # self._calls.append(
        #     _Call(
        #         "declare_channel",
        #         (channel_name, channel_id),
        #         {"initial_target": initial_target},
        #     )
        # )

    def _validate_channel(self, channel_name: str) -> None:
        if channel_name not in self._channels:
            raise ValueError("Use the channel name of a declared channel.")

    def _add_to_schedule(self, channel_name: str, timeslot: _TimeSlot) -> None:
        if hasattr(self, "_measurement"):
            raise RuntimeError(
                "The sequence has already been measured. "
                "Nothing more can be added."
            )
        self._schedule[channel_name].append(timeslot)

    def _delay(self, duration: int, channel: str) -> None:
        self._validate_channel(channel)
        # if self.is_parametrized():
        #     return

        last = self._last(channel)
        ti = last.tf
        tf = ti + self._channels[channel].validate_duration(duration)
        self._add_to_schedule(
            channel, _TimeSlot("delay", ti, tf, last.targets)
        )

    def _last(self, channel: str) -> _TimeSlot:
        """Shortcut to last element in the channel's schedule."""
        try:
            return self._schedule[channel][-1]
        except IndexError:
            raise ValueError("The chosen channel has no target.")

    def _target(
        self, qubits: Union[Iterable[QubitId], QubitId], channel: str
    ) -> None:
        self._validate_channel(channel)

        try:
            qubits_set = (
                set(cast(Iterable, qubits))
                if not isinstance(qubits, str)
                else {qubits}
            )
        except TypeError:
            qubits_set = {qubits}

        if len(qubits_set) > cast(int, self._channels[channel].max_targets):
            raise ValueError(
                "This channel can target at most "
                f"{self._channels[channel].max_targets} qubits at a time"
            )

        if not qubits_set.issubset(self._qids):
            raise ValueError("All given qubits must belong to the register.")

        try:
            last = self._last(channel)
            if last.targets == qubits_set:
                return
            ti = last.tf
            retarget = cast(int, self._channels[channel].retarget_time)
            elapsed = ti - self._last_target[channel]
            delta = cast(int, np.clip(retarget - elapsed, 0, retarget))
            if delta != 0:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    delta = self._channels[channel].validate_duration(
                        16 if delta < 16 else delta
                    )
            tf = ti + delta

        except ValueError:
            ti = -1
            tf = 0

        self._last_target[channel] = tf
        self._add_to_schedule(channel, _TimeSlot("target", ti, tf, qubits_set))

    def _get_time_index(self, t: Union[float, int], t_tol: float = 1.0e-3) -> int:
        duration = self.get_duration()
        sim_times = np.array(range(duration+1))

        diff = np.where(abs(t - sim_times) < t_tol)[0]
        if len(diff)==0:
            raise IndexError(
                f"Given time {t} is absent from Simulation times within"
                + f" tolerance {t_tol}."
            )
        else:
            return diff[0]

    def _min_pulse_duration(self) -> float:
        duration_list = []
        for ch_schedule in self._schedule.values():
            for slot in ch_schedule:
                if isinstance(slot.type, tuple):
                    duration_list.append(slot.tf - slot.ti)
        return min(duration_list)

    def add(
        self,
        Ω: Union[Waveform, float, int],
        Δ: Union[Waveform, float, int],
        ϕ: Union[Waveform, float, int],
        channel_name: str,
        duration: float = None,
        protocol: str = "min-delay",
    ) -> None:

        # Validate the input
        durations = [None, None, None]
        if isinstance(Ω, Waveform): durations[0] = Ω.duration
        if isinstance(Δ, Waveform): durations[1] = Δ.duration
        if isinstance(ϕ, Waveform): durations[2] = ϕ.duration

        durations2 = [x for x in durations if x is not None]

        if len(set(durations2))==0 and duration is None:
            raise ValueError("If the inputs Ω, Δ and ϕ are all numbers, the duration of the waveforms needs to be provided")
        elif len(set(durations2)) >=2:
            raise ValueError("The durations of the input waveforms should be equal to each other")
        else:
            if len(set(durations2))>0: duration = durations2[0]
            if durations[0] == None: Ω = ConstantWaveform(duration, Ω)
            if durations[1] == None: Δ = ConstantWaveform(duration, Δ)
            if durations[2] == None: ϕ = ConstantWaveform(duration, ϕ)

        channel_name = cast(str, channel_name)
        self._validate_channel(channel_name)

        valid_protocols = ["min-delay", "no-delay", "wait-for-all"]
        if protocol not in valid_protocols:
            raise ValueError(
                f"Invalid protocol '{protocol}', only accepts protocols: "
                + ", ".join(valid_protocols)
            )

        last_TimeSlot = self._last(channel_name)
        t0 = last_TimeSlot.tf  # Preliminary ti

        # basis = channel_obj.basis
        # phase_barriers = [
        #     self._phase_ref[basis][q].last_time for q in last_TimeSlot.targets
        # ]
        current_max_t = t0 # max(t0, *phase_barriers)
        if protocol != "no-delay":
            for ch, _ in self._schedule.items():
                if ch == channel_name:
                    continue
                for op in self._schedule[ch][::-1]:
                    if op.tf <= current_max_t:
                        break
                    # if not isinstance(op.type, Pulse):
                    if not isinstance(op.type, Tuple):
                        continue
                    if op.targets & last_TimeSlot.targets or protocol == "wait-for-all":
                        current_max_t = op.tf
                        break

        # figure out ti,
        ti = current_max_t
        if ti > t0:
            # Insert a delay
            delay_duration = ti - t0

            # Delay must not be shorter than the min duration for this channel_name
            min_duration = self._channels[channel_name].min_duration
            if delay_duration < min_duration:
                ti += min_duration - delay_duration
                delay_duration = min_duration

            self._delay(delay_duration, channel_name)

        # figure out tf,
        tf = ti + Ω.duration # pulse.duration

        # prs = {self._phase_ref[basis][q].last_phase for q in last.targets}
        # if len(prs) != 1:
        #     raise ValueError(
        #         "Cannot do a multiple-target pulse on qubits with different "
        #         "phase references for the same basis."
        #     )
        # else:
        #     phase_ref = prs.pop()

        # if phase_ref != 0:
        #     # Has to copy to keep the original pulse intact
        #     pulse = copy.deepcopy(pulse)
        #     pulse.phase = (pulse.phase + phase_ref) % (2 * np.pi)

        # figure out targets
        targets = last_TimeSlot.targets

        self._add_to_schedule(channel_name, _TimeSlot((Ω, Δ, ϕ), ti, tf, targets))

    def cut(self, t_initial: Union[float, int], t_final: Union[float, int], t_tol: float = 1.0e-3) -> Sequence:
        """Returns a portion of the given sequence between the given initial and final time.

        Keyword Args:
            channel (Optional[str]): A specific channel to return the duration
                of. If left as None, it will return the duration of the whole
                sequence.

        Returns:
            Sequence: The duration of the channel or sequence, in ns.
        """
        ind_initial = self._get_time_index(t_initial)
        ind_final = self._get_time_index(t_final)
        duration = self.get_duration()

        if ind_initial >= ind_final:
            raise ValueError(f"The given t_initial {t_initial} is not strictly smaller than the given t_final {t_final} within tolerance {t_tol}")

        if ind_initial==duration:
            pass # This is an error which has been addressed above
        else:
            # seq = copy.deepcopy(self)
            lattice = self._lattice
            seq = Sequence(lattice)
            seq._channels = self.declared_channels

            for ch in seq.declared_channels.keys():
                seq._schedule[ch] = [self._schedule[ch][0]]

                intervals = []
                for timeslot in self._schedule[ch]:
                    intervals.append([timeslot.ti, timeslot.tf])

                # print(ch, intervals)

                if len(intervals)==1:
                    continue
                else:
                    t0 = seq._last(ch).tf
                    for ind, interval in enumerate(intervals[1:]):
                        ind += 1
                        if interval[1]<=ind_initial or interval[0]>=ind_final:
                            pass
                        else:
                            # print("t0", t0)
                            timeslot = self._schedule[ch][ind]
                            timeslot_duration = timeslot.tf-timeslot.ti

                            # Determine the set of points in this interval to keep
                            ti = max(ind_initial, timeslot.ti)
                            tf = min(ind_final, timeslot.tf)

                            if ti==timeslot.ti:
                                ind2 = range(tf-ti)
                            elif tf == timeslot.tf:
                                ind2 = range(ti-(tf-timeslot_duration), timeslot_duration)
                            else:
                                ind2 = range(ti-timeslot.ti, tf-timeslot.ti)

                            # print("ti, tf, timeslot.ti, timeslot.tf, ind2 ", ti, tf, timeslot.ti, timeslot.tf, ind2)

                            # shift the starting point to match the previous interval
                            tf = tf - (ti-t0)
                            ti = t0
                            t0 = tf

                            if type(timeslot.type) == tuple:
                                Ω = CustomWaveform([timeslot.type[0].samples[i] for i in ind2])
                                Δ = CustomWaveform([timeslot.type[1].samples[i] for i in ind2])
                                ϕ = CustomWaveform([timeslot.type[2].samples[i] for i in ind2])
                                new_timeslot = _TimeSlot((Ω, Δ, ϕ), ti, tf, timeslot.targets)
                            else:
                                new_timeslot = _TimeSlot(timeslot.type, ti, tf, timeslot.targets)
                            seq._add_to_schedule(ch, new_timeslot)

            return seq

    def get_duration(self, channel: Optional[str] = None) -> int:
        """Returns the current duration of a channel or the whole sequence.

        Keyword Args:
            channel (Optional[str]): A specific channel to return the duration
                of. If left as None, it will return the duration of the whole
                sequence.

        Returns:
            int: The duration of the channel or sequence, in ns.
        """
        if channel is None:
            durations = [
                self._last(ch).tf
                for ch in self._schedule
                if self._schedule[ch]
            ]
            return 0 if not durations else max(durations)

        self._validate_channel(channel)
        return self._last(channel).tf if self._schedule[channel] else 0

    def draw(
        self,
        draw_phase_area: bool = False,
        # draw_interp_pts: bool = True,
        # draw_phase_shifts: bool = False,
    ) -> None:
        """Draws the sequence in its current state.

        Keyword Args:
            draw_phase_area (bool): Whether phase and area values need to be
                shown as text on the plot, defaults to False.
            draw_interp_pts (bool): When the sequence has pulses with waveforms
                of type InterpolatedWaveform, draws the points of interpolation
                on top of the respective waveforms (defaults to True).
            draw_phase_shifts (bool): Whether phase shift and reference
                information should be added to the plot, defaults to False.

        See Also:
            Simulation.draw(): Draws the provided sequence and the one used by
            the solver.
        """
        draw_sequence(
            self,
            draw_phase_area=draw_phase_area,
            # draw_interp_pts=draw_interp_pts,
            # draw_phase_shifts=draw_phase_shifts,
        )