from copy import copy
from .nmf import NMF
from .make_template import TEMPLATE_PATH, HOP_SIZE, SR
from .make_template import BASIS, FRAME_SIZE, ATTACK, BINS
from .utils import make_pianoroll, find_start_stop, midipath2mat
from .utils import stretch_pianoroll, mat2midipath
import pickle
import torch
import torch.nn.functional as F
from torch import nn
import numpy as np
import sys
import random
from tqdm import tqdm

MINI_SPEC_SIZE = 14
DEVICE = 'cuda'
ALIGNED_MINI_SPEC_PATH = 'aligned_mini_specs.pkl'
VIENNA_MINI_SPEC_PATH = 'vienna_mini_specs.pkl'
ALIGNED_VELOCITY_MODEL_PATH = 'aligned_velocity_model.pkl'
VIENNA_VELOCITY_MODEL_PATH = 'vienna_velocity_model.pkl'
COST_FUNC = 'EucDist'
NJOBS = 5
EPS_ACTIVATIONS = 1e-4
NUM_SONGS_FOR_TRAINING = 100
EPOCHS = 500
BATCH_SIZE = 400
EARLY_STOP = 10
BRANCHES = 16
DATASET_LEN = 1  # use this for debugging
EPS_RANGE = 0.1

def spectrogram(audio, frames=FRAME_SIZE, hop=HOP_SIZE):

    import essentia.standard as esst
    import essentia as es
    spectrogram = []
    spec = esst.SpectrumCQ(numberBins=BINS, sampleRate=SR, windowType='hann')
    for frame in esst.FrameGenerator(audio, frameSize=frames, hopSize=hop):
        spectrogram.append(spec(frame))

    return es.array(spectrogram).T

def get_default_predict_func(vienna_based):
    """
    Return the default  predict function based on PyTorch.  If `vienna_based`
    is True, it will be the model trained with vienna transcription method,
    otherwise it will be the model based on magenta aligned score
    """
    velocity_model = VelocityEstimation().to(DEVICE)
    if vienna_based:
        parameters = pickle.load(open(VIENNA_VELOCITY_MODEL_PATH, 'rb'))
    else:
        parameters = pickle.load(open(ALIGNED_VELOCITY_MODEL_PATH, 'rb'))
    velocity_model.load_state_dict(parameters)
    velocity_model.eval()
    return velocity_model.predict

def transcribe(audio,
               data,
               score=None,
               res=0.001,
               sr=SR,
               return_mini_specs=False):
    """
    Takes an audio mono file and the non-aligned score mat format as in asmd.
    Align them and perform NMF with default templates.
    Returns new score with velocities and timings updated

    `res` is only used for alignment
    """
    if not return_mini_specs:
        velocity_model = get_default_predict_func(score is None)

    initW, minpitch, maxpitch = data
    initW = copy(initW)
    if score is not None:
        from .alignment.align_with_amt import audio_to_score_alignment
        score = copy(score)
        # align score
        new_ons, new_offs = audio_to_score_alignment(score, audio, sr, res=res)
        score[:, 1] = new_ons
        score[:, 2] = new_offs
    else:
        from .vienna_transcription import transcribe
        score = transcribe(audio, sr)

    # prepare initial matrices

    # remove stoping and starting silence in audio
    start, stop = find_start_stop(audio, sample_rate=sr)
    audio = audio[start:stop]
    V = spectrogram(audio)

    # compute the needed resolution for pianoroll
    res = len(audio) / sr / V.shape[1]
    pr = make_pianoroll(score,
                        res=res,
                        basis=BASIS,
                        velocities=False,
                        attack=ATTACK,
                        eps=EPS_ACTIVATIONS,
                        eps_range=EPS_RANGE)

    # remove trailing zeros in initH
    nonzero_cols = pr.any(axis=0).nonzero()[0]
    start = nonzero_cols[0]
    stop = nonzero_cols[-1]
    pr = pr[:, start:stop + 1]

    # stretch pianoroll
    initH = stretch_pianoroll(pr, V.shape[1])

    # check shapes
    assert V.shape == (initW.shape[0], initH.shape[1]),\
        "V, W, H shapes are not comparable"
    assert initH.shape[0] == initW.shape[1],\
        "W, H have different ranks"

    initW = initW[:, minpitch * BASIS:(maxpitch + 1) * BASIS]
    initH = initH[minpitch * BASIS:(maxpitch + 1) * BASIS, :]
    initH[initH == 0] = EPS_ACTIVATIONS

    # perform nfm
    NMF(V, initW, initH, B=BASIS, num_iter=5, cost_func=COST_FUNC)

    NMF(V, initW, initH, B=BASIS, num_iter=5, cost_func=COST_FUNC, fixW=True)

    # use the updated H and W for computing mini-spectrograms
    # and predict velocities
    mini_specs = []
    npitch = maxpitch - minpitch + 1
    initH = initH.reshape(npitch, BASIS, -1)
    initW = initW.reshape((-1, npitch, BASIS), order='C')
    # removing existing velocities
    score[:, 3] = -255
    for note in score:
        # extract mini-spectrogram

        # look for the maximum value in initH in the note
        start = max(0, int(note[1] / res))
        end = min(initH.shape[2], int(note[2] / res))

        if end - start < 1:
            note[3] = 63
            mini_specs.append(None)
            continue

        m = np.argmax(
            np.max(initH[int(note[0] - minpitch), :, start:end],
                   axis=0)) + start

        # select the sorrounding space in initH
        start = max(0, m - MINI_SPEC_SIZE // 2)
        end = min(start + MINI_SPEC_SIZE, initH.shape[2])

        if end - start < MINI_SPEC_SIZE:
            note[3] = 63
            mini_specs.append(None)
            continue

        # compute the mini_spec
        mini_spec = initW[:, int(note[0] - minpitch), :] @\
            initH[int(note[0] - minpitch), :, start:end]

        # normalizing with rms
        # mini_spec /= (mini_spec**2).mean()**0.5
        # normalizing to the sum
        mini_spec /= mini_spec.sum()

        mini_specs.append(mini_spec)

    if return_mini_specs:
        return mini_specs
    else:
        # remove nans...
        mini_specs = [i for i in mini_specs if i is not None]
        # numpy to torch and add channel dimensions
        mini_specs = torch.tensor(mini_specs).to(DEVICE).to(
            torch.float).unsqueeze(1)
        with torch.no_grad():
            vels = velocity_model(mini_specs)
        score[score[:, 3] != 63, 3] = vels.cpu().numpy()
        return score, V, initW, initH

def transcribe_from_paths(audio_path,
                          data,
                          velocity_model,
                          midi_score_path=None,
                          tofile='out.mid'):
    """
    Load a midi and an audio file and call `transcribe`. If `tofile` is not
    empty, it will also write a new MIDI file with the provided path.
    The output midi file will contain only one track with piano (program 0)
    """
    import essentia.standard as esst
    audio = esst.EasyLoader(filename=audio_path, sampleRate=SR)()
    if midi_score_path:
        score = midipath2mat(midi_score_path)
    else:
        score = None
    new_score, _, _, _ = transcribe(audio,
                                    data,
                                    score=score,
                                    velocity_model=velocity_model)

    # writing to midi
    mat2midipath(new_score, tofile)
    return new_score

def processing(i, dataset, data):
    audio, sr = dataset.get_mix(i, sr=SR)
    score = dataset.get_score(i, score_type=['non_aligned'])
    velocities = dataset.get_score(i, score_type=['precise_alignment'])[:, 3]
    return transcribe(audio, data, score=score,
                      return_mini_specs=True), velocities.tolist()

def create_mini_specs(data, mini_spec_path):
    """
    Perform alignment and NMF but not velocity estimation; instead, saves all
    the mini_specs of each note in the Maestro dataset for successive training
    """
    from asmd.asmd import audioscoredataset
    from .maestro_split_indices import maestro_splits
    train, validation, test = maestro_splits()
    dataset = audioscoredataset.Dataset().filter(datasets=["Maestro"])
    random.seed(1750)
    train = random.sample(train, NUM_SONGS_FOR_TRAINING)
    dataset.paths = np.array(dataset.paths)[train].tolist()

    data = dataset.parallel(processing, data, n_jobs=NJOBS)

    mini_specs, velocities = [], []
    for d in data:
        specs, vels = d
        # removing nones
        for i in range(len(specs)):
            spec = specs[i]
            vel = vels[i]
            if spec is not None and vel is not None:
                mini_specs.append(spec)
                velocities.append(vel)

    pickle.dump((mini_specs, velocities), open(mini_spec_path, 'wb'))
    print(
        f"number of (inputs, targets) in training set: {len(mini_specs)}, {len(velocities)}"
    )

class VelocityEstimation(nn.Module):
    def __init__(self,
                 in_numel=MINI_SPEC_SIZE * 100,
                 branches=BRANCHES,
                 k=128 // BRANCHES + 1):
        super().__init__()

        self.preprocess = nn.Sequential(nn.BatchNorm2d(1), nn.Dropout(0.3))

        self.in_numel = in_numel
        self.branches = branches
        self.k = k

        self.process = nn.ModuleList()

        for i in range(branches):
            self.process.append(
                nn.Sequential(nn.Linear(in_numel, k, bias=True), nn.SELU()))

        self.finalize = nn.Sequential(
            nn.Linear(branches * k, branches * k, bias=False), nn.SELU(),
            nn.Linear(branches * k, branches * k, bias=False), nn.SELU(),
            nn.Linear(branches * k, branches * k, bias=False), nn.SELU(),
            nn.Linear(branches * k, branches * k, bias=False), nn.SELU(),
            nn.Linear(branches * k, branches * k, bias=False), nn.SELU(),
            nn.Linear(branches * k, 1, bias=False), nn.Sigmoid())

        # self.apply(lambda x: init_weights(x, nn.init.kaiming_uniform_))

    def forward(self, x):

        # preprocess
        x = self.preprocess(x).reshape(x.shape[0], -1)

        # process each velocity range
        y = torch.zeros(x.shape[0], self.branches,
                        self.k).to(x.dtype).to(x.device)
        for i in range(self.branches):
            y[:, i, :] = self.process[i](x)

        # apply softmax so that only the first output is a probability (classification)
        middle_out = F.softmax(y[:, :, 0], dim=1)

        # finalize takes as input the concatenation of all the features of previous layers
        if self.k > 1:
            x = torch.cat([y[..., i] for i in range(1, self.k)], dim=1)
            x = torch.cat([middle_out, x], dim=1)
        else:
            x = middle_out
        x = self.finalize(x)[:, 0] * 127

        return x, middle_out

    def predict(self, x):
        x = self.forward(x)[0]
        return x
        # return torch.argmax(x, dim=1)

def init_weights(m, initializer):
    if hasattr(m, "weight"):
        if m.weight is not None:

            w = m.weight.data
            if w.dim() < 2:
                w = w.unsqueeze(0)
            initializer(w)

class Dataset(torch.utils.data.Dataset):
    def __init__(self, inputs, targets, branches=BRANCHES):
        super().__init__()
        self.inputs = torch.tensor(inputs).to(torch.float).to(DEVICE)
        self.targets = torch.tensor(targets).to(torch.float).to(DEVICE)
        self.targets_middle = torch.zeros(len(targets),
                                          branches).to(torch.float).to(DEVICE)
        self.targets_middle[torch.arange(len(targets)), targets % branches] = 1
        assert len(self.inputs) == len(self.targets),\
            "inputs and targets must have the same length!"
        del inputs, targets

    def __getitem__(self, i):
        return self.inputs[i], self.targets[i], self.targets_middle[i]

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

def train(data, model_path, mini_spec_path):

    print("Loading dataset...")
    mini_spec = open(mini_spec_path, 'rb')
    inputs, targets = pickle.load(mini_spec)
    mini_spec.close()

    print("Building model...")
    model = VelocityEstimation().to(DEVICE)
    print(model)

    # shuffle and split
    indices = list(range(len(inputs) // DATASET_LEN))
    random.seed(1998)
    random.shuffle(indices)
    inputs = np.array(inputs)
    targets = np.array(targets)
    train_size = int(len(indices) * 0.7)
    test_size = valid_size = int(len(indices) * 0.15)
    train_x = inputs[indices[:train_size]]
    valid_x = inputs[indices[train_size:train_size + valid_size]]
    test_x = inputs[indices[-test_size:]]
    train_y = targets[indices[:train_size]]
    valid_y = targets[indices[train_size:train_size + valid_size]]
    test_y = targets[indices[-test_size:]]

    # creating loaders
    trainloader = torch.utils.data.DataLoader(Dataset(train_x, train_y,
                                                      BRANCHES),
                                              batch_size=BATCH_SIZE)
    validloader = torch.utils.data.DataLoader(Dataset(valid_x, valid_y,
                                                      BRANCHES),
                                              batch_size=BATCH_SIZE)
    testloader = torch.utils.data.DataLoader(Dataset(test_x, test_y, BRANCHES),
                                             batch_size=BATCH_SIZE)
    del train_x, train_y, valid_x, valid_y, test_x, test_y, inputs, targets

    optim = torch.optim.Adadelta(model.parameters(), lr=1e-3)

    best_epoch = 0
    best_params = None
    best_loss = 9999
    for epoch in range(EPOCHS):
        print(f"-- Epoch {epoch} --")
        trainloss, validloss = [], []
        print("-> Training")
        model.train()
        for inputs, targets, targets_middle in tqdm(trainloader):
            inputs = inputs.to(DEVICE).unsqueeze(1)
            targets = targets.to(DEVICE)
            targets_middle = targets_middle.to(DEVICE)

            optim.zero_grad()
            out, middle_out = model(inputs)
            bce_loss = F.binary_cross_entropy(middle_out, targets_middle)
            l1_loss = F.l1_loss(out, targets)
            # loss = l1_loss
            loss = bce_loss + l1_loss
            loss.backward()
            optim.step()
            trainloss.append(l1_loss.detach().cpu().numpy())

        print(f"training loss : {np.mean(trainloss)}")

        print("-> Validating")
        with torch.no_grad():
            model.eval()
            for inputs, targets, _ in tqdm(validloader):
                inputs = inputs.unsqueeze(1)
                targets = targets.to(DEVICE)
                # targets = torch.argmax(targets, dim=1).to(torch.float)

                out = model.predict(inputs).to(torch.float)
                loss = torch.abs(targets - out)
                validloss += loss.tolist()

        validloss = np.mean(validloss)
        print(f"validation loss : {validloss}")
        if validloss < best_loss:
            best_loss = validloss
            best_epoch = epoch
            best_params = model.state_dict()
        elif epoch - best_epoch > EARLY_STOP:
            print("-- Early stop! --")
            break

    # saving params
    model.load_state_dict(best_params)
    pickle.dump(model.to('cpu').state_dict(), open(model_path, 'wb'))
    model.to(DEVICE)

    # testing
    print("-> Testing")
    testloss = []
    with torch.no_grad():
        model.eval()
        for inputs, targets, _ in tqdm(testloader):
            inputs = inputs.unsqueeze(1)
            targets = targets.to(DEVICE)
            # targets = torch.argmax(targets, dim=1).to(torch.float)

            out = model.predict(inputs).to(torch.float)
            loss = torch.abs(targets - out)
            testloss += loss.tolist()

        print(
            f"testing absolute error (mean, std): {np.mean(testloss)}, {np.std(testloss)}"
        )

def show_usage():
    print(
        f"Usage: {sys.argv[0]} [audio_path midi_output_path [midi_score_path] [--cpu]]"
    )
    print(f"Usage: {sys.argv[0]} create_mini_specs, [--vienna]")
    print(f"Usage: {sys.argv[0]} train [--vienna]")

if __name__ == '__main__':
    if len(sys.argv) < 2:
        show_usage()
    elif sys.argv[1] == 'create_mini_specs':

        data = pickle.load(open(TEMPLATE_PATH, 'rb'))
        mini_spec_path = ALIGNED_MINI_SPEC_PATH
        if '--vienna' in sys.argv:
            mini_spec_path = VIENNA_MINI_SPEC_PATH
        create_mini_specs(data, mini_spec_path)

    elif sys.argv[1] == 'train':

        data = pickle.load(open(TEMPLATE_PATH, 'rb'))
        mini_spec_path = ALIGNED_MINI_SPEC_PATH
        model_path = ALIGNED_VELOCITY_MODEL_PATH
        if '--vienna' in sys.argv:
            model_path = VIENNA_VELOCITY_MODEL_PATH
            mini_spec_path = VIENNA_MINI_SPEC_PATH
        train(data, model_path, mini_spec_path)

    elif len(sys.argv) < 3:
        show_usage()
    else:

        data = pickle.load(open(TEMPLATE_PATH, 'rb'))

        if len(sys.argv) > 3:
            if '--cpu' in sys.argv:
                DEVICE = 'cpu'
            else:
                score = sys.argv[3]
        else:
            score = None
        transcribe_from_paths(sys.argv[1],
                              data,
                              midi_score_path=score,
                              tofile=sys.argv[2])


import numpy
import conf
def PRINTER():
	# IMPLICIT #real*8 (A-H,O-Z)
	# IMPLICIT #integer*8 (I-N)
	#integer*4 NSEED
	#COMMON/INPT/
	global NGAS,NSTEP,NANISO,EFINAL,ESTEP,AKT,ARY,TEMPC,TORR,IPEN
	#COMMON/INPT2/
	global KGAS,LGAS,DETEFF,EXCWGHT
	#COMMON/INPT1/
	global NDVEC
	#COMMON/COMP/
	global LCMP,LCFLG,LRAY,LRFLG,LPAP,LPFLG,LBRM,LBFLG,LPEFLG
	#COMMON/RATIO/
	global AN1,AN2,AN3,AN4,AN5,AN6,AN
	global FRAC#(6)
	#COMMON/SETP/
	global TMAX,SMALL,API,ESTART,THETA,PHI
	global TCFMAX#(10),
	global TCFMAX1,RSTART,EFIELD,ETHRM,ECUT,NEVENT,IMIP,IWRITE
	#COMMON/BFLD/
	global EOVB,WB,BTHETA,BMAG
	#COMMON/IONC/
	global DOUBLE#(6,20000),
	global CMINIXSC#(6),
	global CMINEXSC#(6),
	global ECLOSS#(6),
	global WPLN#(6),
	global ICOUNT,AVPFRAC#(3,6)
	#COMMON/LARGE/
	global CF#(20000,512),
	global EIN#(512),
	global TCF#(20000),
	global IARRY#(512),
	global RGAS#(512),
	global IPN#(512),
	global WPL#(512),
	global IZBR#(512),
	global IPLAST,PENFRA#(3,512)
	#COMMON/NAMES/
	global NAMEG#(6)
	#COMMON/KSEED/
	global NSEED
	#COMMON/ECASC/
	global NEGAS#(512),
	global LEGAS#(512),
	global IESHELL#(512),
	global IECASC
	NGAS=conf.NGAS
	NSTEP=conf.NSTEP
	NANISO=conf.NANISO
	EFINAL=conf.EFINAL
	ESTEP=conf.ESTEP
	AKT=conf.AKT
	ARY=conf.ARY
	TEMPC=conf.TEMPC
	TORR=conf.TORR
	IPEN=conf.IPEN
	KGAS=conf.KGAS
	LGAS=conf.LGAS
	DETEFF=conf.DETEFF
	EXCWGHT=conf.EXCWGHT
	NDVEC=conf.NDVEC
	LCMP=conf.LCMP
	LCFLG=conf.LCFLG
	LRAY=conf.LRAY
	LRFLG=conf.LRFLG
	LPAP=conf.LPAP
	LPFLG=conf.LPFLG
	LBRM=conf.LBRM
	LBFLG=conf.LBFLG
	LPEFLG =conf.LPEFLG
	AN1=conf.AN1
	AN2=conf.AN2
	AN3=conf.AN3
	AN4=conf.AN4
	AN5=conf.AN5
	AN6=conf.AN6
	AN=conf.AN
	FRAC=conf.FRAC
	TMAX=conf.TMAX
	SMALL=conf.SMALL
	API=conf.API
	ESTART=conf.ESTART
	THETA=conf.THETA
	PHI=conf.PHI
	TCFMAX=conf.TCFMAX

	TCFMAX1=conf.TCFMAX1
	RSTART=conf.RSTART
	EFIELD=conf.EFIELD
	ETHRM=conf.ETHRM
	ECUT=conf.ECUT
	NEVENT=conf.NEVENT
	IMIP=conf.IMIP
	IWRITE=conf.IWRITE
	EOVB=conf.EOVB
	WB=conf.WB
	BTHETA=conf.BTHETA
	BMAG  =conf.BMAG
	DOUBLE=conf.DOUBLE

	CMINIXSC=conf.CMINIXSC

	CMINEXSC=conf.CMINEXSC

	ECLOSS=conf.ECLOSS

	WPLN=conf.WPLN

	ICOUNT=conf.ICOUNT
	AVPFRAC=conf.AVPFRAC
	CF=conf.CF

	EIN=conf.EIN

	TCF=conf.TCF

	IARRY=conf.IARRY

	RGAS=conf.RGAS

	IPN=conf.IPN

	WPL=conf.WPL

	IZBR=conf.IZBR

	IPLAST=conf.IPLAST
	PENFRA=conf.PENFRA
	NAMEG=conf.NAMEG
	NSEED =conf.NSEED
	NEGAS=conf.NEGAS

	LEGAS=conf.LEGAS

	IESHELL=conf.IESHELL

	IECASC  =conf.IECASC
	# NAMEG=numpy.zeros(25+1,dtype=str)
	# WRITE(6,1)
	print('\n           DEGRAD VERSION 3.3  \n','      -----------------------------\n\n')
	if(IMIP == 1):
		print('   MIP AND DE/DX SIMULATION')	#2
	if(IMIP == 2):
		print('   ELECTRON BEAM SIMULATION')
	if(IMIP == 3):
		print('   X-RAY SIMULATION')	#4
	if(IMIP == 4):
		print('   BETA DECAY SIMULATION')	#5
	if(IMIP == 5):
		print('   DOUBLE BETA DECAY SIMULATION')	#6
	print('----------------------------------\n\n')
	if(LCMP == 0):
		print('   SIMULATION WITHOUT COMPTON SCATTERING')  	#7
	if(LCMP == 1):
		print('   SIMULATION WITH COMPTON SCATTERING')	#8
	if(LRAY == 0):
		print('   SIMULATION WITHOUT RAYLEIGH SCATTERING')	#9
	if(LRAY == 1):
		print('   SIMULATION WITH RAYLEIGH SCATTERING')	#11
	if(LPAP == 0):
		print('   SIMULATION WITHOUT PAIR PRODUCTION')	#12
	if(LPAP == 1):
		print('   SIMULATION WITH PAIR PRODUCTION')	#13
	if(LBRM == 0):
		print('   SIMULATION WITHOUT BREMSSTRAHLUNG')	#14
	if(LBRM == 1):
		print('   SIMULATION WITH BREMSSTRAHLUNG')	#15
	if(IECASC == 0):
		print('   SIMULATION WITH PARAMETERISED SHELL CASCADE')	#16
	if(IECASC == 1):
		print('   SIMULATION WITH COMPLETE SHELL CASCADE')	#17
	print('----------------------------------\n\n')
	print('   MONTE CARLO SOLUTION FOR MIXTURE OF ',NGAS,' GASES.\n   DEGRADATION CALCULATION ALL TIMES IN PICOSECS, DISTANCE IN MICRONS\n   -----------------------------------------------------------------')
	# WRITE(6,30) (NAMEG[J],FRAC[J], J=1,NGAS)
	for J in range(1,NGAS+1):
		print('\n',5*' ','  GASES  USED ',15*' ',' PERCENTAGE USED ',2*'\n',6*' ',NAMEG[J],5*' ','%.4f' % FRAC[J],'\n')
	print('\n','  ','GAS TEMPERATURE =','%.1f' % TEMPC,' DEGREES CENTIGRADE.','\n','  ','GAS PRESSURE = ','%.1f' % TORR,' TORR.')
	if(NSEED != 0):
		# WRITE(6,51) NSEED
		print(2*'\n',' RANDOM NUMBER SEED =',NSEED)
	if(NSEED == 0):
		# WRITE(6,52)
		print(2*'\n',' STANDARD RANDOM NUMBER SEED = 54217137')
	if(IPEN == 0):
		# WRITE(6,55)
		print(2*'\n','  ',' PENNING IONISATION NOT ALLOWED')
	if(IPEN == 1):
		# WRITE(6,56)
		print(2*'\n','  ',' PENNING IONISATION ALLOWED')
	# WRITE(6,60) EFINAL,NSTEP
	print(1*'\n','  ','INTEGRATION FROM 0.0 TO ','%.1f' % EFINAL,' EV.  IN ',NSTEP,' STEPS. ')
	# WRITE(6,90) EFIELD,BMAG,BTHETA,WB
	print(1*'\n','  ELECTRIC FIELD =','%.4f' % EFIELD,' VOLTS/CM.','\n''  MAGNETIC FIELD =','%.4f' % BMAG,' KILOGAUSS.','\n','  ANGLE BETWEEN ELECTRIC AND MAGNETIC FIELD =','%.3f' % BTHETA,' DEGREES.','\n','  CYCLOTRON FREQ. =','%.3f' % WB,' RADIANS/PICOSECOND')
	# WRITE(6,43)
	print('\n',' USED ANISOTROPIC X-SECTIONS (OKHRIMOVSKYY ET AL) ')
	if(ICOUNT == 1):
		# WRITE(6,34)
		print(' USED COUNTING IONISATION X-SECTIONS')
	else:
		# WRITE(6,35)
		print(' USE GROSS IONISATION X-SECTIONS')
	# endif
	# WRITE(6,91) ESTART,NEVENT,ETHRM
	# print(NEVENT,conf.NEVENT)
	print(1*'\n','  INITIAL ELECTRON OR X-RAY ENERGY =','%.1f' % ESTART,' EV.','\n',9*' ','NUMBER OF EVENTS =',NEVENT,'\n',4*' ','THERMALISATION ENERGY =','%.2f' % ETHRM,' EV.','\n')
	# WRITE(6,911) DETEFF,EXCWGHT
	print(' PHOTON DETECTION EFFICIENCY USED IN FANO CALCULATION =','%.3f' % DETEFF,' %','\n',7*' ','WEIGHT GIVEN TO EXCITATION IN FANO CALCULATION =','%.3f' % EXCWGHT,'\n')
	# print(IMIP)
	if(IMIP == 4 or IMIP == 5):
		if(KGAS <= 0 or KGAS > NGAS):
			# WRITE(6,990) KGAS
			print(' ERROR IN INPUT: BETA DECAY IDENTifIER KGAS=',KGAS,'  PROGRAM STOPPED:')
			sys.exit()
		# endif
		if(LGAS <= 0 or LGAS > 3):
			# WRITE(6,991) LGAS
			print(' ERROR IN INPUT: BETA DECAY IDENTIFIER LGAS=',LGAS,'  PROGRAM STOPPED:')
			sys.exit()
		# endif
		# WRITE(6,88) KGAS,LGAS
		print('\n  BETA DECAY IN GAS NO =',KGAS,'\n  IF MOLECULE : BETA DECAY IN ATOMIC COMPONENT =',LGAS,'\n')
	# endif
	if(NDVEC == 2):
		# WRITE(6,915)
		print('  BETA OR X-RAY IN RANDOM DIRECTION TO E-FIELD')
		pass
	else:
		# endif
		if(abs(numpy.cos(THETA)) < 1.e-9 and IMIP > 2):
			# WRITE(6,92)
			print('  BETA OR X-RAY PERP# endICULAR TO E-FIELD IN X-Y PLANE')
		if(abs(numpy.cos(THETA)) < 1.e-9 and IMIP == 2):
			# WRITE(6,922)
			print('  ELECTRON BEAM ALONG X DIRECTION')
		if(numpy.cos(THETA)== 1.0):
			# WRITE(6,93)
			print('  E-BEAM,BETA OR X-RAY ALONG Z-AXIS IN E-FIELD DIRECTION')
		if(numpy.cos(THETA)== -1.0):
			# WRITE(6,94)
			print('  E-BEAM,BETA OR X-RAY ALONG Z-AXIS OPPOSITE TO E-FIELD DIRECTION')
	# 95  WRITE(6,96) TCFMAX1

	print("TCFMAX1",TCFMAX1,type(TCFMAX1))
	print('\n NULL COLLISION FREQUENCY = %.4f *(10**12/SEC)\n'%(TCFMAX1))
	# WRITE(6,111)  (TCF(L),L=500,9500,1000)
	print('  ','REAL COLLISION FREQUENCY AT 10 EQUALLY SPACED ENERGY INTERVALS (*10**12/SEC)','\n')
	for L in range(500,9500+1,1000):
		print(3*' ','%.3f' % TCF[L],'\t', end='')
		if L==4500:
			print('\n')
	print('\n')

	conf.NGAS=NGAS
	conf.NSTEP=NSTEP
	conf.NANISO=NANISO
	conf.EFINAL=EFINAL
	conf.ESTEP=ESTEP
	conf.AKT=AKT
	conf.ARY=ARY
	conf.TEMPC=TEMPC
	conf.TORR=TORR
	conf.IPEN=IPEN
	conf.KGAS=KGAS
	conf.LGAS=LGAS
	conf.DETEFF=DETEFF
	conf.EXCWGHT=EXCWGHT
	conf.NDVEC=NDVEC
	conf.LCMP=LCMP
	conf.LCFLG=LCFLG
	conf.LRAY=LRAY
	conf.LRFLG=LRFLG
	conf.LPAP=LPAP
	conf.LPFLG=LPFLG
	conf.LBRM=LBRM
	conf.LBFLG=LBFLG
	conf.LPEFLG =LPEFLG
	conf.AN1=AN1
	conf.AN2=AN2
	conf.AN3=AN3
	conf.AN4=AN4
	conf.AN5=AN5
	conf.AN6=AN6
	conf.AN=AN
	conf.FRAC=FRAC
	conf.TMAX=TMAX
	conf.SMALL=SMALL
	conf.API=API
	conf.ESTART=ESTART
	conf.THETA=THETA
	conf.PHI=PHI
	conf.TCFMAX=TCFMAX

	conf.TCFMAX1=TCFMAX1
	conf.RSTART=RSTART
	conf.EFIELD=EFIELD
	conf.ETHRM=ETHRM
	conf.ECUT=ECUT
	conf.NEVENT=NEVENT
	conf.IMIP=IMIP
	conf.IWRITE=IWRITE
	conf.EOVB=EOVB
	conf.WB=WB
	conf.BTHETA=BTHETA
	conf.BMAG  =BMAG
	conf.DOUBLE=DOUBLE

	conf.CMINIXSC=CMINIXSC

	conf.CMINEXSC=CMINEXSC

	conf.ECLOSS=ECLOSS

	conf.WPLN=WPLN

	conf.ICOUNT=ICOUNT
	conf.AVPFRAC=AVPFRAC
	conf.CF=CF

	conf.EIN=EIN

	conf.TCF=TCF

	conf.IARRY=IARRY

	conf.RGAS=RGAS

	conf.IPN=IPN

	conf.WPL=WPL

	conf.IZBR=IZBR

	conf.IPLAST=IPLAST
	conf.PENFRA=PENFRA
	conf.NSEED =NSEED
	conf.NEGAS=NEGAS

	conf.LEGAS=LEGAS

	conf.IESHELL=IESHELL

	conf.IECASC  =IECASC
	return
	# end


# 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.
"""All built-in types of waveforms and their Waveform parent class."""

from __future__ import annotations

from abc import ABC, abstractmethod
import functools
import inspect
import itertools
import sys
from sys import version_info
from types import FunctionType
from typing import Any, cast, Optional, Tuple, Union
import warnings

from matplotlib.axes import Axes
import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import ArrayLike
import scipy.interpolate as interpolate

# from pulser.parametrized import Parametrized, ParamObj
# from pulser.parametrized.decorators import parametrize
# from pulser.json.utils import obj_to_dict

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

class Waveform(ABC):
    """The abstract class for a pulse's waveform."""

    # def __new__(cls, *args, **kwargs):  # type: ignore
    #     """Creates a Waveform instance or a ParamObj depending on the input."""
    #     for x in itertools.chain(args, kwargs.values()):
    #         if isinstance(x, Parametrized):
    #             return ParamObj(cls, *args, **kwargs)
    #     else:
    #         return object.__new__(cls)

    def __init__(self, duration: int): # (self, duration: Union[int, Parametrized]):
        """Initializes a waveform with a given duration.

        Args:
            duration (int): The waveforms duration (in ns).
        """
        duration = cast(int, duration)
        try:
            _duration = int(duration)
        except (TypeError, ValueError):
            raise TypeError(
                "duration needs to be castable to an int but "
                f"type {type(duration)} was provided."
            )
        if _duration <= 0:
            raise ValueError(
                "A waveform must have a positive duration, "
                + f"not {duration}."
            )
        elif duration - _duration != 0:
            warnings.warn(
                f"A waveform duration of {duration} ns is below the"
                " supported precision of 1 ns. It was rounded down "
                + f"to {_duration} ns.",
                stacklevel=3,
            )

        self._duration = _duration

    @property
    @abstractmethod
    def duration(self) -> int:
        """The duration of the pulse (in ns)."""
        pass

    @cached_property
    @abstractmethod
    def _samples(self) -> np.ndarray:
        pass

    @property
    def samples(self) -> np.ndarray:
        """The value at each time step that describes the waveform.

        Returns:
            np.ndarray: A numpy array with a value for each time step.
        """
        return self._samples.copy()

    @property
    def first_value(self) -> float:
        """The first value in the waveform."""
        return float(self[0])

    @property
    def last_value(self) -> float:
        """The last value in the waveform."""
        return float(self[-1])

    @property
    def integral(self) -> float:
        """Integral of the waveform (time in ns, value in rad/µs)."""
        return float(np.sum(self.samples)) * 1e-3  # ns * rad/µs = 1e-3

    def draw(self) -> None:
        """Draws the waveform."""
        fig, ax = plt.subplots()
        self._plot(ax, "rad/µs")

        plt.show()

    def change_duration(self, new_duration: int) -> Waveform:
        """Returns a new waveform with modified duration.

        Args:
            new_duration(int): The duration of the new waveform.
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} does not support"
            " modifications to its duration."
        )

    # @abstractmethod
    # def _to_dict(self) -> dict[str, Any]:
    #     pass

    @abstractmethod
    def __str__(self) -> str:
        pass

    @abstractmethod
    def __repr__(self) -> str:
        pass

    def __getitem__(
        self, index_or_slice: Union[int, slice]
    ) -> Union[float, np.ndarray]:
        if isinstance(index_or_slice, slice):
            s: slice = self._check_slice(index_or_slice)
            return cast(np.ndarray, self._samples[s])
        else:
            index: int = self._check_index(index_or_slice)
            return cast(float, self._samples[index])

    def _check_index(self, i: int) -> int:
        if i < -self.duration or i >= self.duration:
            raise IndexError(
                "Index ('index_or_slice' = "
                f"{i}) must be in the range "
                f"0~{self.duration-1}, or "
                f"{-self.duration}~-1 from the end."
            )
        return i if i >= 0 else self.duration + i

    def _check_slice(self, s: slice) -> slice:
        if s.step is not None and s.step != 1:
            raise IndexError("The step of the slice must be None or 1.")

        # Transform start and stop indexes into positive or null values
        # since they can be omitted (None) or negative (end-indexing)
        start = (
            0
            if s.start is None
            else (s.start if s.start >= 0 else self.duration + s.start)
        )
        stop = (
            self.duration
            if s.stop is None
            else (s.stop if s.stop >= 0 else self.duration + s.stop)
        )

        # Correct out of bounds ranges
        if start < 0:
            start = 0
        if stop < 0:
            stop = 0
        if start > self.duration:
            start = self.duration
        if stop > self.duration:
            stop = self.duration
        if stop < start:
            stop = start

        return slice(start, stop)

    @abstractmethod
    def __mul__(self, other: float) -> Waveform:
        pass

    def __neg__(self) -> Waveform:
        return self.__mul__(-1.0)

    def __truediv__(self, other: float) -> Waveform:
        if other == 0:
            raise ZeroDivisionError("Can't divide a waveform by zero.")
        else:
            return self.__mul__(1 / other)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Waveform):
            return False
        elif self.duration != other.duration:
            return False
        else:
            return bool(np.all(np.isclose(self.samples, other.samples)))

    def __hash__(self) -> int:
        return hash(tuple(self.samples))

    def _plot(
        self, ax: Axes, ylabel: str, color: Optional[str] = None
    ) -> None:
        ax.set_xlabel("t (ns)")
        ts = np.arange(self.duration)
        if color:
            ax.set_ylabel(ylabel, color=color, fontsize=14)
            ax.plot(ts, self.samples, color=color)
            ax.tick_params(axis="y", labelcolor=color)
            ax.axhline(0, color=color, linestyle=":", linewidth=0.5)
        else:
            ax.set_ylabel(ylabel, fontsize=14)
            ax.plot(ts, self.samples, 'o-')
            ax.axhline(0, color="black", linestyle=":", linewidth=0.5)

class LinearWaveform(Waveform):
    """A linear ramp waveform.

    Args:
        duration (int): The waveform duration (in ns).
        start (float): The initial value (in rad/µs).
        stop (float): The final value (in rad/µs).
    """

    def __init__(
        self,
        duration: int, # Union[int, Parametrized],
        start: float, # Union[float, Parametrized],
        stop: float, # Union[float, Parametrized],
    ):
        """Initializes a linear waveform."""
        super().__init__(duration)
        start = cast(float, start)
        self._start: float = float(start)
        stop = cast(float, stop)
        self._stop: float = float(stop)

    @property
    def duration(self) -> int:
        """The duration of the pulse (in ns)."""
        return self._duration

    @cached_property
    def _samples(self) -> np.ndarray:
        """The value at each time step that describes the waveform.

        Returns:
            numpy.ndarray: A numpy array with a value for each time step.
        """
        return np.linspace(self._start, self._stop, num=self._duration)

    @property
    def slope(self) -> float:
        r"""Slope of the ramp, in :math:`s^{-15}`."""
        return (self._stop - self._start) / self._duration

    def change_duration(self, new_duration: int) -> LinearWaveform:
        """Returns a new waveform with modified duration.

        Args:
            new_duration(int): The duration of the new waveform.

        Returns:
            LinearWaveform: The new waveform with the given duration.
        """
        return LinearWaveform(new_duration, self._start, self._stop)

    # def _to_dict(self) -> dict[str, Any]:
    #     return obj_to_dict(self, self._duration, self._start, self._stop)

    def __str__(self) -> str:
        return f"LinearWaveform({self._start:.3g}->{self._stop:.3g} rad/µs)"

    def __repr__(self) -> str:
        return (
            f"LinearWaveform({self._duration} ns, "
            + f"{self._start:.3g}->{self._stop:.3g} rad/µs)"
        )

    def __mul__(self, other: float) -> LinearWaveform:
        k = float(other)
        return LinearWaveform(self._duration, self._start * k, self._stop * k)

class ConstantWaveform(Waveform):
    """A waveform of constant value.

    Args:
        duration (int): The waveform duration (in ns).
        value (float): The modulation value (in rad/µs).
    """

    def __init__(
        self,
        duration: int, # Union[int, Parametrized],
        value: float, # Union[float, Parametrized],
    ):
        """Initializes a constant waveform."""
        super().__init__(duration)
        value = cast(float, value)
        self._value = float(value)

    @property
    def duration(self) -> int:
        """The duration of the pulse (in ns)."""
        return self._duration

    @cached_property
    def _samples(self) -> np.ndarray:
        """The value at each time step that describes the waveform.

        Returns:
            numpy.ndarray: A numpy array with a value for each time step.
        """
        return np.full(self.duration, self._value)

    def change_duration(self, new_duration: int) -> ConstantWaveform:
        """Returns a new waveform with modified duration.

        Args:
            new_duration(int): The duration of the new waveform.

        Returns:
            ConstantWaveform: The new waveform with the given duration.
        """
        return ConstantWaveform(new_duration, self._value)

    # def _to_dict(self) -> dict[str, Any]:
    #     return obj_to_dict(self, self._duration, self._value)

    def __str__(self) -> str:
        return (
            f"ConstantWaveform({self._duration} ns, "
            + f"{self._value:.3g} rad/µs)"
        )

    def __repr__(self) -> str:
        return (
            f"ConstantWaveform({self._duration} ns, "
            + f"{self._value:.3g} rad/µs)"
        )

    def __mul__(self, other: float) -> ConstantWaveform:
        return ConstantWaveform(self._duration, self._value * float(other))

class CompositeWaveform(Waveform):
    """A waveform combining multiple smaller waveforms.

    Args:
        waveforms(Waveform): Two or more waveforms to combine.
    """

    def __init__(self, *waveforms: Waveform):
        """Initializes a waveform from multiple waveforms."""
        if len(waveforms) < 2:
            raise ValueError(
                "Needs at least two waveforms to form a " "CompositeWaveform."
            )
        waveforms = cast(Tuple[Waveform], waveforms)
        for wf in waveforms:
            self._validate(wf)

        self._waveforms = list(waveforms)

    @property
    def duration(self) -> int:
        """The duration of the pulse (in ns)."""
        duration = 0
        for wf in self._waveforms:
            duration += wf.duration
        return duration

    @cached_property
    def _samples(self) -> np.ndarray:
        """The value at each time step that describes the waveform.

        Returns:
            numpy.ndarray: A numpy array with a value for each time step.
        """
        return cast(
            np.ndarray, np.concatenate([wf.samples for wf in self._waveforms])
        )

    @property
    def waveforms(self) -> list[Waveform]:
        """The waveforms encapsulated in the composite waveform."""
        return list(self._waveforms)

    def _validate(self, waveform: Waveform) -> None:
        if not isinstance(waveform, Waveform):
            raise TypeError(
                f"{waveform!r} is not a valid waveform. "
                "Please provide a valid Waveform."
            )

    # def _to_dict(self) -> dict[str, Any]:
    #     return obj_to_dict(self, *self._waveforms)

    def __str__(self) -> str:
        contents_list = ["{!r}"] * len(self._waveforms)
        contents = ", ".join(contents_list)
        contents = contents.format(*self._waveforms)
        return f"Composite({contents})"

    def __repr__(self) -> str:
        return f"CompositeWaveform({self.duration} ns, {self._waveforms!r})"

    def __mul__(self, other: float) -> CompositeWaveform:
        return CompositeWaveform(*(wf * other for wf in self._waveforms))

class CustomWaveform(Waveform):
    """A custom waveform.

    Args:
        samples (array_like): The modulation values at each time step
            (in rad/µs). The number of samples dictates the duration, in ns.
    """

    def __init__(self, samples: ArrayLike):
        """Initializes a custom waveform."""
        samples_arr = np.array(samples, dtype=float)
        self._samples: np.ndarray = samples_arr
        super().__init__(len(samples_arr))

    @property
    def duration(self) -> int:
        """The duration of the pulse (in ns)."""
        return self._duration

    @cached_property
    def _samples(self) -> np.ndarray:
        """The value at each time step that describes the waveform.

        Returns:
            numpy.ndarray: A numpy array with a value for each time step.
        """
        # self._samples is already cached when initialized in __init__
        pass

    # def _to_dict(self) -> dict[str, Any]:
    #     return obj_to_dict(self, self._samples)

    def __str__(self) -> str:
        return "Custom"

    def __repr__(self) -> str:
        return f"CustomWaveform({self.duration} ns, {self.samples!r})"

    def __mul__(self, other: float) -> CustomWaveform:
        return CustomWaveform(self._samples * float(other))

class InterpolatedWaveform(Waveform):
    """Creates a waveform from interpolation of a set of data points.

    Args:
        duration (int): The waveform duration (in ns).
        values (ArrayLike): Values of the interpolation points (in rad/µs).
        times (Optional[ArrayLike]): Fractions of the total duration (between 0
            and 1), indicating where to place each value on the time axis. If
            not given, the values are spread evenly throughout the full
            duration of the waveform.
        interpolator (str = "PchipInterpolator"): The SciPy interpolation class
            to use. Supports "PchipInterpolator" and "interp1d".
        **interpolator_kwargs: Extra parameters to give to the chosen
            interpolator class.
    """

    def __init__(
        self,
        duration: int, # Union[int, Parametrized],
        values: ArrayLike, # Union[ArrayLike, Parametrized],
        times: Optional[ArrayLike] = None, # Optional[Union[ArrayLike, Parametrized]] = None,
        interpolator: str = "PchipInterpolator",
        **interpolator_kwargs: Any,
    ):
        """Initializes a new InterpolatedWaveform."""
        super().__init__(duration)
        self._values = np.array(values, dtype=float)
        if times is not None:
            times_ = np.array(times, dtype=float)
            if len(times_) != len(self._values):
                raise ValueError(
                    "When specified, the number of time coordinates in `times`"
                    f" ({len(times_)}) must match the number of `values` "
                    f"({len(self._values)})."
                )
            if np.any(times_ < 0):
                raise ValueError(
                    "All values in `times` must be greater than or equal to 0."
                )
            if np.any(times_ > 1):
                raise ValueError(
                    "All values in `times` must be less than or equal to 1."
                )
            unique_times = np.unique(times)  # Sorted array of unique values
            if len(times_) != len(unique_times):
                raise ValueError(
                    "`times` must be an array of non-repeating values."
                )
            self._times = times_
        else:
            self._times = np.linspace(0, 1, num=len(self._values))

        valid_interpolators = ("PchipInterpolator", "interp1d")
        if interpolator not in valid_interpolators:
            raise ValueError(
                f"Invalid interpolator '{interpolator}', only "
                "accepts: " + ", ".join(valid_interpolators)
            )
        interp_cls = getattr(interpolate, interpolator)
        self._data_pts = np.array(
            [
                (round(t), v)
                for t, v in zip(
                    self._times * (self._duration - 1), self._values
                )
            ]
        )
        self._interp_func = interp_cls(
            self._data_pts[:, 0], self._data_pts[:, 1], **interpolator_kwargs
        )
        self._kwargs: dict[str, Any] = {
            "times": times,
            "interpolator": interpolator,
            **interpolator_kwargs,
        }

    @property
    def duration(self) -> int:
        """The duration of the pulse (in ns)."""
        return self._duration

    @cached_property
    def _samples(self) -> np.ndarray:
        """The value at each time step that describes the waveform."""
        return cast(
            np.ndarray,
            np.round(
                self._interp_func(np.arange(self._duration)), decimals=9
            ),  # Rounds to the order of Hz
        )

    @property
    def interp_function(
        self,
    ) -> Union[interpolate.PchipInterpolator, interpolate.interp1d]:
        """The interpolating function."""
        return self._interp_func

    @property
    def data_points(self) -> np.ndarray:
        """Points (t[ns], value[rad/µs]) that define the interpolation."""
        return self._data_pts.copy()

    def change_duration(self, new_duration: int) -> InterpolatedWaveform:
        """Returns a new waveform with modified duration.

        Args:
            new_duration(int): The duration of the new waveform.

        Returns:
            InterpolatedWaveform: The new waveform with the same coordinates
            for interpolation but a new duration.
        """
        return InterpolatedWaveform(new_duration, self._values, **self._kwargs)

    def _plot(
        self, ax: Axes, ylabel: str, color: Optional[str] = None
    ) -> None:
        super()._plot(ax, ylabel, color=color)
        ax.scatter(self._data_pts[:, 0], self._data_pts[:, 1], c=color)

    # def _to_dict(self) -> dict[str, Any]:
    #     return obj_to_dict(self, self._duration, self._values, **self._kwargs)

    def __str__(self) -> str:
        coords = [f"({int(x)}, {y:.4g})" for x, y in self.data_points]
        return f"InterpolatedWaveform(Points: {', '.join(coords)})"

    def __repr__(self) -> str:
        interp_str = f", Interpolator={self._kwargs['interpolator']})"
        return self.__str__()[:-1] + interp_str

    def __mul__(self, other: float) -> InterpolatedWaveform:
        return InterpolatedWaveform(
            self._duration, self._values * other, **self._kwargs
        )

# # To replicate __init__'s signature in __new__ for every Waveform subclass
# def _copy_func(f: FunctionType) -> FunctionType:
#     return FunctionType(
#         f.__code__,
#         f.__globals__,
#         name=f.__name__,
#         argdefs=f.__defaults__,
#         closure=f.__closure__,
#     )

# for m in inspect.getmembers(sys.modules[__name__], inspect.isclass):
#     if m[1].__module__ == __name__:
#         _new = _copy_func(m[1].__new__)
#         m[1].__new__ = functools.update_wrapper(_new, m[1].__init__)


"""
# @Time    : 2021/7/2 5:22 下午
# @Author  : hezhiqiang01
# @Email   : hezhiqiang01@baidu.com
# @File    : env.py
"""
#只需要编写这一部分的代码，就可以无缝衔接MAPPO。
import numpy as np

import math
import gym
from gym import spaces, logger
from gym.utils import seeding

class Env(object):
    """
    # 环境中的智能体
    """
    def __init__(self, i):
        # self.agent_num = 3  # 设置智能体(小飞机)的个数，这里设置为两个
        # self.obs_dim = 4  # 设置智能体的观测纬度14
        # self.action_dim = 2  # 设置智能体的动作纬度，这里假定为一个五个纬度的

       #CartPoleEnv

        # # Angle limit set to 2 * theta_threshold_radians so failing observation
        self.agent_num = 1  # 设置智能体的个数，这里设置为两个
        self.obs_dim = 4  # 设置智能体的观测纬度
        self.action_dim = 2  # 设置智能体的动作纬度，这里假定为一个五个纬度的

        self.gravity = 9.8
        self.masscart = 1.0
        self.masspole = 0.1
        self.total_mass = (self.masspole + self.masscart)
        self.length = 0.5  # actually half the pole's length
        self.polemass_length = (self.masspole * self.length)
        self.force_mag = 10.0
        self.tau = 0.02  # seconds between state updates
        self.kinematics_integrator = 'euler'  # 运动学积分仪

        # Angle at which to fail the episode
        self.theta_threshold_radians = 12 * 2 * math.pi / 360
        self.x_threshold = 2.4
        # is still within bounds.
        high = np.array([self.x_threshold * 2,
                         np.finfo(np.float32).max,
                         self.theta_threshold_radians * 2,
                         np.finfo(np.float32).max],
                        dtype=np.float32)

        self.action_space = spaces.Discrete(2)  # env.step(0) ：小车向左，env.step(1) ：小车向右.定义了一个变量空间范围为[0,2) 之间的整数
        self.observation_space = spaces.Box(-high, high, dtype=np.float32)  # ,义了一个取值范围在（-10，10）的变量 维度为1

        self.seed()
        self.viewer = None
        self.state = None

        self.steps_beyond_done = None

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def reset(self):
        """
        # self.agent_num设定为2个智能体时，返回值为一个list，每个list里面为一个shape = (self.obs_dim, )的观测数据
        """
        ##new
        self.state = self.np_random.uniform(low=-0.05, high=0.05, size=(4,))
        self.steps_beyond_done = None
        return np.array(self.state)

        sub_agent_obs = []
        for i in range(self.agent_num):
            sub_obs = np.random.random(size=(self.obs_dim, ))#14，obs_dim
            sub_agent_obs.append(sub_obs)
        return sub_agent_obs

        #new
        # sub_agent_obs = []
        # for i in range(self.agent_num):
        #     self.state = self.np_random.uniform(low=-0.05, high=0.05, size=(4,))
        #     sub_agent_obs.append(self.state)
        # return sub_agent_obs

# # if true, action is a number 0...N, otherwise action is a one-hot N-dimensional vector
        #self.discrete_action_input = False
    def step(self, actions):
        """
        # self.agent_num设定为2个智能体时，actions的输入为一个2纬的list，每个list里面为一个shape = (self.action_dim, )的动作数据
        # 默认参数情况下，输入为一个list，里面含有两个元素，因为动作纬度为5，所里每个元素shape = (5, )
        """
        # sub_agent_obs = []
        # sub_agent_reward = []
        # sub_agent_done = []
        # sub_agent_info = []
        # for i in range(self.agent_num):
        #     sub_agent_obs.append(np.random.random(size=(15,)))
        #     sub_agent_reward.append([np.random.rand()])
        #     sub_agent_done.append(False)
        #     sub_agent_info.append({})
        #
        # #print('sub_agent_obs',sub_agent_obs, 'sub_agent_reward', sub_agent_reward, 'sub_agent_done', sub_agent_done, 'sub_agent_info',sub_agent_info)
        # return [sub_agent_obs, sub_agent_reward, sub_agent_done, sub_agent_info]

    #new
        #err_msg = "%r (%s) invalid" % (actions, type(actions))
        #assert self.action_space.contains(actions), err_msg
        sub_agent_obs = []
        sub_agent_reward = []
        sub_agent_done = []
        sub_agent_info = []
        for i in range(self.agent_num):
            x, x_dot, theta, theta_dot = self.state
            #print("actions[i]",actions)
            force = self.force_mag if actions[i][0] == 1 else -self.force_mag
            # if action == 1:
            #     force = self.force_mag
            # elif action == 0:
            #     force = -self.force_mag
            # else:
            #     force = 0
            costheta = math.cos(theta)
            sintheta = math.sin(theta)

            # For the interested reader:
            # https://coneural.org/florian/papers/05_cart_pole.pdf
            temp = (force + self.polemass_length * theta_dot ** 2 * sintheta) / self.total_mass
            thetaacc = (self.gravity * sintheta - costheta * temp) / (
                        self.length * (4.0 / 3.0 - self.masspole * costheta ** 2 / self.total_mass))
            xacc = temp - self.polemass_length * thetaacc * costheta / self.total_mass

            if self.kinematics_integrator == 'euler':
                x = x + self.tau * x_dot
                x_dot = x_dot + self.tau * xacc
                theta = theta + self.tau * theta_dot
                theta_dot = theta_dot + self.tau * thetaacc
            else:  # semi-implicit euler
                x_dot = x_dot + self.tau * xacc
                x = x + self.tau * x_dot
                theta_dot = theta_dot + self.tau * thetaacc
                theta = theta + self.tau * theta_dot

            self.state = (x, x_dot, theta, theta_dot)  # (位置x，x加速度, 偏移角度theta, 角加速度)
            # 小车的世界，就一条x轴，
            # 变量env.x_threshold里存放着小车坐标的最大值（=2.4），
            # 超过这个数值，世界结束，每step()一次，就会奖励 1，直到上次done为True。
            done = bool(
                x < -self.x_threshold
                or x > self.x_threshold
                or theta < -self.theta_threshold_radians
                or theta > self.theta_threshold_radians
            )

            if not done:
                reward = 1.0
            elif self.steps_beyond_done is None:
                # Pole just fell!
                self.steps_beyond_done = 0
                reward = 1.0
            else:
                if self.steps_beyond_done == 0:
                    logger.warn(
                        "You are calling 'step()' even though this "
                        "environment has already returned done = True. You "
                        "should always call 'reset()' once you receive 'done = "
                        "True' -- any further steps are undefined behavior."
                    )
                self.steps_beyond_done += 1
                reward = 0.0
            sub_agent_obs.append(np.array(self.state))
            sub_agent_reward.append(reward)
            sub_agent_done.append(done)
            sub_agent_info.append({})

        return [sub_agent_obs, sub_agent_reward, sub_agent_done, sub_agent_info]


# --*-- coding:utf-8 --*--
import numpy as np
import cv2
from scipy import signal

'''
helper function
'''
def filterItChopOff(f, r, sp):
    f[np.isnan(f)] = 0
    H, W, d = f.shape
    B = np.ones([2 * r + 1, 2 * r + 1])     # 2r+1 * 2r+1 neighbourhood

    minSP = cv2.erode(sp, B, iterations=1)
    maxSP = cv2.dilate(sp, B, iterations=1)

    ind = np.where(np.logical_or(minSP != sp, maxSP != sp))

    spInd = np.reshape(range(np.size(sp)), sp.shape,'F')

    delta = np.zeros(f.shape)
    delta = np.reshape(delta, (H * W, d), 'F')
    f = np.reshape(f, (H * W, d),'F')

    # calculate delta

    I, J = np.unravel_index(ind, [H, W], 'C')
    for i in range(np.size(ind)):
        x = I[i]
        y = J[i]
        clipInd = spInd[max(0, x - r):min(H-1, x + r), max(0, y - r):min(W-1, y + r)]
        diffInd = clipInd[sp[clipInd] != sp[x, y]]
        delta[ind[i], :] = np.sum(f[diffInd, :], 1)
    delta = np.reshape(delta, (H, W, d), 'F')
    f = np.reshape(f, (H, W, d), 'F')
    fFilt = np.zeros([H, W, d])

    for i in range(f.shape[2]):
        #  fFilt(:,:,i) = filter2(B, f(:,:,i));
        tmp = signal.convolve2d(np.rot90(f[:, :, i], 2), np.rot90(np.rot90(B, 2), 2), mode="same")
        fFilt[:, :, i] = np.rot90(tmp, 2)
    fFilt = fFilt - delta
    return fFilt

'''
helper function
'''
def mutiplyIt(AtA_1, Atb):
    result = np.zeros([Atb.shape[0], Atb.shape[1], 3])
    result[:, :, 0] = np.multiply(AtA_1[:, :, 0], Atb[:, :, 0]) + np.multiply(AtA_1[:, :, 1],
                                                                              Atb[:, :, 1]) + np.multiply(
        AtA_1[:, :, 2], Atb[:, :, 2])
    result[:, :, 1] = np.multiply(AtA_1[:, :, 1], Atb[:, :, 0]) + np.multiply(AtA_1[:, :, 3],
                                                                              Atb[:, :, 1]) + np.multiply(
        AtA_1[:, :, 4], Atb[:, :, 2])
    result[:, :, 2] = np.multiply(AtA_1[:, :, 2], Atb[:, :, 0]) + np.multiply(AtA_1[:, :, 4],
                                                                              Atb[:, :, 1]) + np.multiply(
        AtA_1[:, :, 5], Atb[:, :, 2])
    return result

'''
helper function
'''
def invertIt(AtA):
    AtA_1 = np.zeros([AtA.shape[0], AtA.shape[1], 6])
    AtA_1[:, :, 0] = np.multiply(AtA[:, :, 3], AtA[:, :, 5]) - np.multiply(AtA[:, :, 4], AtA[:, :, 4])
    AtA_1[:, :, 1] = -np.multiply(AtA[:, :, 1], AtA[:, :, 5]) + np.multiply(AtA[:, :, 2], AtA[:, :, 4])
    AtA_1[:, :, 2] = np.multiply(AtA[:, :, 1], AtA[:, :, 4]) - np.multiply(AtA[:, :, 2], AtA[:, :, 3])
    AtA_1[:, :, 3] = np.multiply(AtA[:, :, 0], AtA[:, :, 5]) - np.multiply(AtA[:, :, 2], AtA[:, :, 2])
    AtA_1[:, :, 4] = -np.multiply(AtA[:, :, 0], AtA[:, :, 4]) + np.multiply(AtA[:, :, 1], AtA[:, :, 2])
    AtA_1[:, :, 5] = np.multiply(AtA[:, :, 0], AtA[:, :, 3]) - np.multiply(AtA[:, :, 1], AtA[:, :, 1])

    x1 = np.multiply(AtA[:, :, 0], AtA_1[:, :, 0])
    x2 = np.multiply(AtA[:, :, 1], AtA_1[:, :, 1])
    x3 = np.multiply(AtA[:, :, 2], AtA_1[:, :, 2])

    detAta = x1 + x2 + x3
    return AtA_1, detAta

'''
Compute the direction of gravity
N: normal field
iter: number of 'big' iterations
'''
def getYDir(N, angleThresh, iter, y0):
    y = y0
    for i in range(len(angleThresh)):
        thresh = np.pi * angleThresh[i] / 180   # convert it to radian measure
        y = getYDirHelper(N, y, thresh, iter[i])
    return y

'''
N: HxWx3 matrix with normal at each pixel.
y0: the initial gravity direction
thresh: in degrees the threshold for mapping to parallel to gravity and perpendicular to gravity
iter: number of iterations to perform
'''
def getYDirHelper(N, y0, thresh, num_iter):
    dim = N.shape[0] * N.shape[1]

    # change the third dimension to the first-order. (480, 680, 3) => (3, 480, 680)
    nn = np.swapaxes(np.swapaxes(N,0,2),1,2)
    nn = np.reshape(nn, (3, dim), 'F')

    # remove these whose number is NAN
    idx = np.where(np.invert(np.isnan(nn[0,:])))[0]
    nn = nn[:,idx]

    # Set it up as a optimization problem
    yDir = y0;
    for i in range(num_iter):
        sim0 = np.dot(yDir.T, nn)
        indF = abs(sim0) > np.cos(thresh)       # calculate 'floor' set.    |sin(theta)| < sin(thresh) ==> |cos(theta)| > cos(thresh)
        indW = abs(sim0) < np.sin(thresh)       # calculate 'wall' set.
        if(len(indF.shape) == 2):
            NF = nn[:, indF[0,:]]
            NW = nn[:, indW[0,:]]
        else:
            NF = nn[:, indF]
            NW = nn[:, indW]
        A = np.dot(NW, NW.T) - np.dot(NF, NF.T)
        b = np.zeros([3,1])
        c = NF.shape[1]
        w,v = np.linalg.eig(A)      # w:eigenvalues; v:eigenvectors
        min_ind = np.argmin(w)      # min index
        newYDir = v[:,min_ind]
        yDir = newYDir * np.sign(np.dot(yDir.T, newYDir))
    return yDir

'''
getRMatrix: Generate a rotation matrix that
            if yf is a scalar, rotates about axis yi by yf degrees
            if yf is an axis, rotates yi to yf in the direction given by yi x yf
Input: yi is an axis 3x1 vector
       yf could be a scalar of axis

'''
def getRMatrix(yi, yf):
    if (np.isscalar(yf)):
        ax = yi / np.linalg.norm(yi)        # norm(A) = max(svd(A))
        phi = yf
    else:
        yi = yi / np.linalg.norm(yi)
        yf = yf / np.linalg.norm(yf)
        ax = np.cross(yi.T, yf.T).T
        ax = ax / np.linalg.norm(ax)
        # find angle of rotation
        phi = np.degrees(np.arccos(np.dot(yi.T, yf)))

    if (abs(phi) > 0.1):
        phi = phi * (np.pi / 180)

        s_hat = np.array([[0, -ax[2], ax[1]],
                          [ax[2], 0, -ax[0]],
                          [-ax[1], ax[0], 0]])
        R = np.eye(3) + np.sin(phi) * s_hat + (1 - np.cos(phi)) * np.dot(s_hat, s_hat)      # dot???
    else:
        R = np.eye(3)
    return R

'''
Calibration of gravity direction
'''
def rotatePC(pc, R):
    if(np.array_equal(R, np.eye(3))):
        return pc
    else:
        dim = pc.shape[0] * pc.shape[1]
        pc = np.swapaxes(np.swapaxes(pc, 0, 2), 1, 2)
        res = np.reshape(pc, (3, dim), 'F')
        res = np.dot(R, res)
        res = np.reshape(res, pc.shape, 'F')
        res = np.swapaxes(np.swapaxes(res, 0, 1), 1, 2)
        return res


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2018 GEM Foundation
#
# OpenQuake 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.
#
# OpenQuake 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 OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module exports :class:'GarciaEtAl2005SSlab',
:class:'GarciaEtAl2005SSlabVert'

"""
import numpy as np
# standard acceleration of gravity in m/s**2
from scipy.constants import g

from openquake.hazardlib.gsim.base import GMPE, CoeffsTable
from openquake.hazardlib import const
from openquake.hazardlib.imt import SA, PGA, PGV

class GarciaEtAl2005SSlab(GMPE):
    """
    Implements GMPE developed by Garcia, D., Singh, S. K., Harraiz, M,
    Ordaz, M., and Pacheco, J. F. and published in BSSA as:

    "Inslab earthquakes of Central Mexico: Peak ground-motion parameters and
    response spectra", vol. 95, No. 6, pp. 2272-2282."

    The original formulation predict peak ground acceleration (PGA), in
    cm/s*s, peak ground velocity PGV (cm/s) and 5% damped pseudo-acceleration
    response spectra (PSA) in cm/s*s for the geometric average of the
    maximum component of the two horizontal component of ground motion (see
    last paragraph of Summary in pag. 2272

    The GMPE predicted values for Mexican inslab events and NEHRP B site
    condition

    """

    #: Supported tectonic region type is subduction intraslab,
    #: given that the equations have been derived using Mexican inslab
    #: events
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.SUBDUCTION_INTRASLAB

    #: Supported intensity measure types are spectral acceleration,
    #: and peak ground acceleration. See Table 2 in page 1865
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        SA,
        PGA,
        PGV,
    ])

    #: Supported intensity measure component is the geometric average of
    #  the maximum of the two horizontal components
    #: :attr:`openquake.hazardlib.const.IMC.AVERAGE_HORIZONTAL`,
    #: see Data processing in page 2274.
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.AVERAGE_HORIZONTAL

    #: Supported standard deviation types are inter-event, intra-event
    #: and total
    #: See Tables 2 and 3, page 2275.
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL,
        const.StdDev.INTER_EVENT,
        const.StdDev.INTRA_EVENT
    ])

    #: No site parameters required
    #: All data from 51 hard (NEHRP B) sites.
    REQUIRES_SITES_PARAMETERS = set(('vs30', ))

    #: Required rupture parameters are magnitude and focal depth
    #: See equation (1) in pag 2274
    REQUIRES_RUPTURE_PARAMETERS = set(('mag', 'hypo_depth'))

    #: Required distance measure is Rrup (closest distance to fault surface for
    #: the larger events, Mw > 6.5) or Rhypo (hypocentral distance for the
    #: rest (both in kilometers) as explained in page 2274
    REQUIRES_DISTANCES = set(('rrup', 'rhypo', ))

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        # Extracting dictionary of coefficients specific to required
        # intensity measure type.

        C = self.COEFFS[imt]
        mag = rup.mag
        hypo_depth = rup.hypo_depth

        mean = self._compute_mean(C, g, mag, hypo_depth, dists, imt)
        stddevs = self._get_stddevs(C, stddev_types, sites.vs30.shape[0])

        return mean, stddevs

    def _compute_mean(self, C, g, mag, hypo_depth, dists, imt):
        """
        Compute mean according to equation on Table 2, page 2275.
        """

        delta = 0.00750 * 10 ** (0.507 * mag)

        # computing R for different values of mag
        if mag < 6.5:
            R = np.sqrt(dists.rhypo ** 2 + delta ** 2)
        else:
            R = np.sqrt(dists.rrup ** 2 + delta ** 2)

        mean = (
            # 1st term
            C['c1'] + C['c2'] * mag +
            # 2nd term
            C['c3'] * R -
            # 3rd term
            C['c4'] * np.log10(R) +
            # 4th term
            C['c5'] * hypo_depth
        )
        # convert from base 10 to base e
        if imt == PGV():
            mean = np.log(10 ** mean)
        else:
            # convert from cm/s**2 to g
            mean = np.log((10 ** mean) * 1e-2 / g)
        return mean

    def _get_stddevs(self, C, stddev_types, num_sites):
        """
        Return standard deviations as defined in table 2, pag 2275.
        """

        assert all(stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
                   for stddev_type in stddev_types)

        # the standard deviation values are converted from base 10 to base e

        stddevs = []
        for stddev_type in stddev_types:
            assert stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
            if stddev_type == const.StdDev.TOTAL:
                stddevs.append(np.log(10 ** C['s_t']) + np.zeros(num_sites))
            elif stddev_type == const.StdDev.INTRA_EVENT:
                stddevs.append(np.log(10 ** C['s_r']) + np.zeros(num_sites))
            elif stddev_type == const.StdDev.INTER_EVENT:
                stddevs.append(np.log(10 ** C['s_e']) + np.zeros(num_sites))
        return stddevs

    #: Equation coefficients for geometric average of the maximum of the two
    #: horizontal components, as described in Table 2 on pp. 2275, but
    #: generated from a Fortran implementation code provided by Daniel Garcia
    #: (higher precision than in the paper).
    #: The original IMT values are defined as frequencies values.
    COEFFS = CoeffsTable(sa_damping=5, table="""\
    IMT   c1       c2        c3       c4   c5       s_t      s_r      s_e
    0.04  0.02645  0.58792  -0.00430  1.0  0.00700  0.31829  0.30777  0.08115
    0.05  0.10949  0.58046  -0.00433  1.0  0.00753  0.33560  0.32329  0.09005
    0.07  0.22907  0.56961  -0.00429  1.0  0.00826  0.33640  0.32126  0.09979
    0.10  0.40746  0.54939  -0.00414  1.0  0.00774  0.33431  0.31785  0.10358
    0.20  0.05215  0.58676  -0.00369  1.0  0.00689  0.27971  0.24464  0.13559
    0.30 -0.26507  0.62932  -0.00331  1.0  0.00485  0.27649  0.22833  0.15592
    0.40 -0.55235  0.64414  -0.00280  1.0  0.00483  0.27187  0.23607  0.13484
    0.50 -0.81731  0.67453  -0.00243  1.0  0.00351  0.26432  0.24081  0.10898
    0.75 -1.31580  0.70924  -0.00198  1.0  0.00371  0.27422  0.25957  0.08843
    1.00 -1.75050  0.75555  -0.00168  1.0  0.00296  0.27728  0.26232  0.08985
    1.50 -2.30120  0.80760  -0.00144  1.0  0.00167  0.28030  0.26085  0.10261
    2.00 -2.75190  0.84564  -0.00123  1.0  0.00137  0.26353  0.24282  0.10240
    3.00 -3.34700  0.89255  -0.00092  1.0  0.00085  0.26279  0.22360  0.13806
    4.00 -3.87460  0.93748  -0.00079  1.0  0.00093  0.25328  0.22226  0.12147
    5.00 -4.26750  0.96929  -0.00074  1.0  0.00104  0.24643  0.21638  0.11793
    pga  -0.23170  0.58726  -0.00394  1.0  0.00767  0.28520  0.26662  0.10123
    pgv  -2.35950  0.70759  -0.00235  1.0  0.00436  0.25745  0.23917  0.09529
    """)

class GarciaEtAl2005SSlabVert(GarciaEtAl2005SSlab):
    """
    Extend :class:`GarciaEtAl2005SSlab`

    Implements GMPE developed by Garcia, D., Singh, S. K., Harraiz, M,
    Ordaz, M., and Pacheco, J. F. and published in BSSA as:

    "Inslab earthquakes of Central Mexico: Peak ground-motion parameters and r
    esponse spectra", vol. 95, No. 6, pp. 2272-2282."

    The original formulation predict peak ground acceleration (PGA), in
    cm/s*s, peak ground velocity PGV (cm/s) and 5% damped pseudo-acceleration
    response spectra (PSA) in cm/s*s for the vertical component of ground
    motion (see last paragraph of Summary in pag. 2272

    The GMPE predicted values for Mexican inslab events and NEHRP B site

    """

    #: Equation coefficients for Vertical Component, as described in Table 3
    #: on pp 2275.
    #: The original imt values are defined as frequencies values
    COEFFS = CoeffsTable(sa_damping=5, table="""\
     IMT    c1     c2      c3      c4     c5      s_t   s_r   s_e
     0.04  -0.300  0.620  -0.0041  1.00   0.0060  0.31  0.30  0.07
     0.05  -0.200  0.620  -0.0043  1.00   0.0070  0.32  0.31  0.08
     0.07  -0.060  0.600  -0.0041  1.00   0.0070  0.32  0.31  0.09
     0.10  -0.040  0.590  -0.0039  1.00   0.0070  0.31  0.29  0.11
     0.20  -0.070  0.590  -0.0033  1.00   0.0040  0.26  0.22  0.14
     0.30  -0.200  0.600  -0.0029  1.00   0.0030  0.26  0.22  0.15
     0.40  -0.700  0.640  -0.0022  1.00   0.0030  0.26  0.23  0.13
     0.50  -0.900  0.660  -0.0018  1.00   0.0020  0.26  0.23  0.11
     0.75  -1.300  0.690  -0.0014  1.00   0.0020  0.25  0.22  0.11
     1.00  -1.800  0.750  -0.0010  1.00   0.0010  0.27  0.24  0.12
     1.50  -2.400  0.800  -0.0008  1.00   0.0004  0.26  0.23  0.12
     2.00  -2.800  0.830  -0.0006  1.00  -0.0005  0.27  0.24  0.14
     3.00  -3.300  0.880  -0.0005  1.00  -0.0004  0.28  0.23  0.17
     4.00  -4.000  0.950  -0.0004  1.00  -0.0003  0.27  0.23  0.15
     5.00  -4.400  0.980  -0.0003  1.00  -0.0002  0.26  0.22  0.14
     pga   -0.400  0.600  -0.0036  1.00   0.0060  0.27  0.25  0.11
     pgv   -2.400  0.700  -0.0018  1.00   0.0020  0.24  0.21  0.11
    """)


# -*- coding: utf-8 -*-
import time
import numpy as np
from pyodesys import ODESys as _ODESys
from pyodesys.results import Result
from chempy.units import get_derived_unit, unitless_in_registry, uniform, patched_numpy as pnp
from .integrate import run
from ._chemreac import cvode_predefined_durations_fields

class ODESys(_ODESys):

    def __init__(self, rd, k_from_params=None, variables_from_params=None):
        if rd.N > 1:
            raise NotImplementedError("ODESys expects single bin for now")
        self.rd = rd
        self.k_from_params = k_from_params
        self.variables_from_params = variables_from_params

    ny = property(lambda self: self.rd.n*self.rd.N)
    names = property(lambda self: self.rd.substance_names)
    latex_names = property(lambda self: self.rd.substance_latex_names)
    param_names = property(lambda self: self.rd.param_names)
    autonomous_interface = property(lambda self: not self.rd.logt)
    numpy = pnp
    # dep_by_name = True
    # par_by_name = True

    def _get_units_util(self):
        if self.rd.unit_registry is None:
            _dedim = lambda x: np.array(x)
            time_u = 1
            conc_u = 1
            dr_u = 1
        else:
            _dedim = lambda x: unitless_in_registry(x, self.rd.unit_registry)
            time_u = get_derived_unit(self.rd.unit_registry, 'time')
            conc_u = get_derived_unit(self.rd.unit_registry, 'concentration')
            dr_u = get_derived_unit(self.rd.unit_registry, 'doserate')
        return locals()

    def integrate(self, x, y0, params=None, integrator='cvode', **kwargs):
        if params is not None and self.k_from_params is not None:
            self.rd.k = self.k_from_params(self, params)
        if 'doserate' in (params or {}):
            self.rd.set_with_units(
                'fields', [[self.variables_from_params['density'](self, params)*params['doserate']]])
        if 'atol' in kwargs and isinstance(kwargs['atol'], dict):
            kwargs['atol'] = [kwargs['atol'][k] for k in self.names]
        integr = run(self.rd, [y0[k] for k in self.names] if isinstance(y0, dict) else y0,
                     x, integrator=integrator, **kwargs)
        pout = [params[k] for k in self.param_names] if self.param_names else None
        return Result(integr.with_units('tout'), integr.with_units('Cout')[:, 0, :],
                      pout, integr.info, self)

    def chained_parameter_variation(self, durations, y0, varied_params, default_params=None,
                                    integrate_kwargs=None, x0=None, npoints=1, numpy=None):
        if list(varied_params) != ['doserate']:
            raise NotImplementedError("For now only varied doserate is supported")
        if self.param_names != ['doserate']:
            raise NotImplementedError("We expect doserate to be varied for now")
        uutil = self._get_units_util()
        _dedim, time_u, conc_u, dr_u = [uutil[k] for k in '_dedim time_u conc_u dr_u'.split()]
        density = _dedim(self.variables_from_params['density'](self, default_params))
        if default_params:
            self.rd.k = _dedim(self.k_from_params(self, default_params))

        if x0 is not None:
            assert x0 == 0*time_u
        integrate_kwargs = integrate_kwargs or {}
        atol = integrate_kwargs.pop('atol', 1e-8)
        if isinstance(atol, float):
            atol = [atol]
        elif isinstance(atol, dict):
            atol = [atol[k] for k in self.names]
        rtol = integrate_kwargs.pop('rtol', 1e-8)
        method = integrate_kwargs.pop('method', 'bdf')
        integrator = integrate_kwargs.pop('integrator', 'cvode')
        if integrator != 'cvode':
            raise NotImplementedError("chained_parameter_variation requires cvode for now")
        drate = uniform(varied_params['doserate'])
        time_cpu = time.process_time()
        time_wall = time.time()
        tout, yout = cvode_predefined_durations_fields(
            self.rd, _dedim([y0[k] for k in self.names]),
            _dedim(durations),
            _dedim(drate*density),
            atol=atol, rtol=rtol, method=method, npoints=npoints, **integrate_kwargs)
        info = dict(
            nsteps=-1,
            nfev=self.rd.nfev,
            njev=self.rd.njev,
            time_wall=time.time() - time_wall,
            time_cpu=time.process_time() - time_cpu,
            success=True,
            integrator=[integrator],
            t0_set=False,
            linear_solver=0,  # pyodesys.results.Result work-around for now (not important)
        )
        info.update(self.rd.last_integration_info)
        dr_out = np.concatenate((np.repeat(drate, npoints), drate[-1:]))
        return Result(tout*time_u, yout[:, 0, :]*conc_u, dr_out.reshape((-1, 1))*dr_u, info, self)


from datetime import datetime
import numpy as np

# equatorial, polar radii
lightning_ellipse_rev = {
    # Values at launch
    0: (6.394140e6, 6.362755e6),

    # DO.07, late 2018. First Virts revision.
    # The GRS80 altitude + 6 km differs by about 3 m from the value above
    # which is the exact that was provided at the time of launch. Use the
    # original value instead of doing the math.
    # 6.35675231414e6+6.0e3
    1: (6.378137e6 + 14.0e3, 6.362755e6),
}

def ltg_ellpse_rev(date):
    """
    Given a date, return the lightning ellipsoid revision. The ellipsoid was
    tuned after launch based on comparison of the GLM flash centroids to ground
    strike locations.

    The date here refers to the operational GLM feed, and this is really
    a convenience function for inferring the lightning ellipsoid revision. An
    operational use of glmtools that expects data from a certain source should
    treat the lightning ellipse revision as a configuration parameter that is
    passed directly to the GLMDataset object (or the gridder class, which passes
    the ellipse version along to GLMDataset). For future revisions of the
    lightning ellipse, this allows for the processing code to be upgraded as
    soon as the new ellipse parameter is known, and the new ellipse to be
    chosen in a live environment as a less-surgery-requiring parameter change.

    date: datetime object for the date and time of observation

    Returns:
    ellps_rev (int): integer used to index the lightning_ellipse_rev dict.

    """
    if date < datetime(2018,10,15):
        return 0
    else:
        return 1

def ltg_ellps_radii(date):
    """
    Given a date, return the equatorial and polar radii of the lightning
    ellipsoid. The ellipsoid was tuned after launch based on comparison of the
    GLM flash centroids to ground strike locations.

    The date here refers to the operational GLM feed, and this is really
    a convenience function for inferring the lightning ellipsoid revision. An
    operational use of glmtools that expects data from a certain source should
    treat the lightning ellipse revision as a configuration parameter that is
    passed directly to the GLMDataset object (or the gridder class, which passes
    the ellipse version along to GLMDataset). For future revisions of the
    lightning ellipse, this allows for the processing code to be upgraded as
    soon as the new ellipse parameter is known, and the new ellipse to be
    chosen in a live environment as a less-surgery-requiring parameter change.

    date: datetime object for the date and time of observation

    Returns: re_ltg_ellps, rp_ltg_ellps (meters): equatorial and polar radii
        of the lightning ellipsoid, respectively.
    """
    re_ltg_ellps, rp_ltg_ellps = lightning_ellipse_rev[ltg_ellps_rev(date)]
    return re_ltg_ellps, rp_ltg_ellps

def ltg_ellps_lon_lat_to_fixed_grid(lon, lat, sat_lon, ellipse_rev,
        re_grs80 = 6.378137e6, rp_grs80 = 6.35675231414e6,
        sat_grs80_height=35.786023e6 ):
    """
    lon, lat (degrees): from GLM L2 file, fixed grid coords x, y as
    defined in the L1b PUG. x,y corresponds to beta, alpha.

    sat_lon (degrees): nominal nadir longitude of the satellite

    re_ltg_ellps, rp_ltg_ellps (meters): equatorial and polar radii
        for the lightning ellipse. Defaults to values set at launch
        of GOES-16 GLM (good at least through early Jan 2018).

    sat_grs80_height (meters): height of the satellite above the GRS80
        ellipsoid. This is 'perspective_point_height' in the GOES-R L1b PUG,
        and the attribute of the same name in the goes_imager_projection
        variable.

    re_grs80, rp_grs80 (meters): equatorial and polar radii
        for GRS80 lightning ellipse used by GOES-R, as defined in the
        GOES-R L1b PUG, and the semi_major_axis and semi_minor_axis
        attributes of the goes_imager_projection variable.

    This function undoes the lightning ellipsoid height assumption,
    such that the final fixed grid position matches the ABI fixed
    grid definition and therefore the ABI L1b products.

    Reference:
    Bezooijen, R. W. H., H. Demroff, G. Burton, D. Chu, and S. Yang, 2016:
        Image navigation and registration for the geostationary lightning
        mapper (GLM). Proc. SPIE 10004, 100041N, doi: 10.1117/12.2242141.
    """
    re_ltg_ellps, rp_ltg_ellps = lightning_ellipse_rev[ellipse_rev]

    ff_ltg_ellps = (re_ltg_ellps - rp_ltg_ellps)/re_ltg_ellps
    ff_grs80 = (re_grs80 - rp_grs80)/re_grs80 # 0.003352810704800
    sat_H = sat_grs80_height + re_grs80 # 42.164e6

    # center longitudes on satellite, and ensure between +/- 180
    lon = np.atleast_1d(lon)
    lat = np.atleast_1d(lat)
    dlon = lon-sat_lon
    dlon[dlon < -180] += 360
    dlon[dlon > 180] -= 360
    lon_rad = np.radians(dlon)
    lat_rad = np.radians(lat)

    lat_geocent = np.arctan( (1.0 - ff_grs80)**2.0 * np.tan(lat_rad))

    # We assume geocentric latitude
    cos_factor = np.cos(lat_geocent)
    sin_factor = np.sin(lat_geocent)

    R = re_ltg_ellps*(1-ff_ltg_ellps) / np.sqrt(1.0-ff_ltg_ellps*(2.0-ff_ltg_ellps)*cos_factor*cos_factor)
    vx = R * cos_factor * np.cos(lon_rad) - sat_H
    vy = R * cos_factor * np.sin(lon_rad)
    vz = R * sin_factor
    vmag = np.sqrt(vx*vx + vy*vy + vz*vz)
    vx /= -vmag # minus signs flip so x points to earth, z up, y left
    vy /= -vmag
    vz /= vmag

    # Microradians
    alpha = np.arctan(vz/vx) #* 1e6
    beta = -np.arcsin(vy) #* 1e6
    return beta, alpha