import numpy as npy
from typing import Dict
import scipy.special

import PyPWA as pwa
from dataclasses import dataclass
import pandas

## Fitting  ##

class FitAmplitude(pwa.NestedFunction):
    
    def __init__(self, initial_params):
        self.__elm = self.make_elm(initial_params)
    
    def setup(self, kVars):
        self.__kVars = kVars
        self.__decay = self.__produce_specific_decay() 
        
    @staticmethod
    def make_elm(params):
        # Number of parameters
        length = int(len(params) / 2)
            
        # Produce ELM Dataframe
        elm = npy.empty(length, [("e", int), ("l", int), ("m", int)])
        elm = pandas.DataFrame(elm)
        
        # Extract ELM values from minuit parameters
        values = []
        for param in params:
            if "r" in param:
                values.append([int(value) for value in param.split(".")[1:]])
            
        # Put extracted values in ELM Dataframe
        for index, value in enumerate(values):
            elm.loc[index] = value
            
        return elm
        
    def __produce_specific_decay(self):
#        thetaC = self.__kVars["phi"].copy()
#        phiC = self.__kVars["theta"].copy()
#        thetaC[thetaC<0] = thetaC[thetaC<0]+2*npy.pi
        complex_amp = npy.empty(len(self.__kVars), "c16")
        decay = npy.empty([len(self.__elm), len(self.__kVars)], "c16")
        for index, wave in self.__elm.iterrows():

#            decay[index] = scipy.special.sph_harm(wave["m"], wave["l"], theta, phi)
#            decay[index] -= wave["e"]*npy.power((-1),wave["m"])*scipy.special.sph_harm(
#                -wave["m"], wave["l"], theta, phi
#            )
        
#        decay = npy.empty([len(wave_data), len(self.__kVars)], "c16")
#       for index, wave in enumerate(wave_data):
            if wave["l"] == 0:
                complex_amp.real= self.__kVars["SAmpx"].copy()
                complex_amp.imag = self.__kVars["SAmpy"].copy()
            elif wave["l"] == 1:
                complex_amp.real= self.__kVars["PAmpx"].copy()
                complex_amp.imag = self.__kVars["PAmpy"].copy()
            elif wave["l"] == 2:
                complex_amp.real= self.__kVars["DAmpx"].copy()
                complex_amp.imag = self.__kVars["DAmpy"].copy()
                
            decay[index] = complex_amp
        return decay
 
    def calculate(self, params):
        intensity = npy.zeros(len(self.__kVars), "c16")
        rho = npy.empty(len(self.__kVars), "c16")
        for wave1 in range(len(self.__elm)):
            for wave2 in range(len(self.__elm)):
                rho.real = 0
                rho.imag = 0
                
                Vs = self.__param_to_Vs(params, wave1, wave2)
                decay = self.__decay[wave1] * self.__decay[wave2].conjugate()
                
                if self.__elm["e"][wave1] == 1 and self.__elm["e"][wave2] == 1:
                    rho.real = 1-0.4*npy.cos(2*self.__kVars["alpha"])
                if self.__elm["e"][wave1] == 1 and self.__elm["e"][wave2] == -1:
                    rho.imag = -1 * 0.4*npy.sin(2 * self.__kVars["alpha"]) 
                if self.__elm["e"][wave1] == -1 and self.__elm["e"][wave2] == 1:
                    rho.imag = 0.4 * npy.sin(2 * self.__kVars["alpha"])
                if self.__elm["e"][wave1] == -1 and self.__elm["e"][wave2] == -1:
                    rho.real = 1 + 0.4 * npy.cos(2 * self.__kVars["alpha"])

#               intensity += .5 * Vs * decay * rho
# no considering polarization

                intensity += Vs * decay
    
#        if npy.sum(intensity.imag):
#            return npy.NaN
           # Kill's the intensity if there is still imaginary values left, and should display
        # those imaginary values
        if npy.sum(intensity.imag) > 1.e-6:
            print("Results are wrong! Imaginary values have been found!")
        self.error = intensity         
        return intensity.real
    
#    def _spin_density(self):
#        
#        npy.empty((2,2),dtype='c16')
#        rho[0][0] = npy.complex(1-self.__kVars["pol"]*npy.cos(2*self.__kVars["alpha"],0.)
#        rho[0][1] = npy.complex(0.,-1. * self.__kVars["pol"]*npy.sin(2 * self.__kVars["alpha"])) 
#       rho[1][0] = npy.complex(0,self.__kVars["pol"] * npy.sin(2 * self.__kVars["alpha"]))
#       rho[1][1] = npy.complex(1 + self.__kVars["pol"] * npy.cos(2 * self.__kVars["alpha"]),0.)
        
#        return 0.5*rho

    def calculate_wave(self, params, wave_num):
        intensity = npy.zeros(len(self.__kVars), "c16")
        rho = npy.empty(len(self.__kVars), "c16")
        wave1 = wave_num

        rho.real = 0
        rho.imag = 0
                
#        Vs = self.__param_to_Vs(params, wave1, wave1)
        real_string = f"r.{self.__elm['e'][wave1]}.{self.__elm['l'][wave1]}.{self.__elm['m'][wave1]}"
        imag_string = f"i.{self.__elm['e'][wave1]}.{self.__elm['l'][wave1]}.{self.__elm['m'][wave1]}"
        Vsc = npy.complex(params[real_string], params[imag_string])
        Vs = Vsc * Vsc.conjugate()
        decay = self.__decay[wave1] * self.__decay[wave1].conjugate()
                
        if self.__elm["e"][wave1] == 1:
            rho.real = 1-0.4*npy.cos(2*self.__kVars["alpha"])
        if self.__elm["e"][wave1] == -1:
            rho.real = 1 + 0.4 * npy.cos(2 * self.__kVars["alpha"])

#       intensity += .5 * Vs * decay * rho
# no considering polarization

        intensity = Vs * decay

        if npy.sum(intensity.imag):
            return npy.NaN
            
        return intensity.real

    def __param_to_Vs(self, params, wave1, wave2):
        value = npy.empty(2, "c16")
        
        # value[0] == wave1; value[1] == wave2
        for index, i in enumerate([wave1, wave2]):
            real_string = f"r.{self.__elm['e'][i]}.{self.__elm['l'][i]}.{self.__elm['m'][i]}"
            imag_string = f"i.{self.__elm['e'][i]}.{self.__elm['l'][i]}.{self.__elm['m'][i]}"
            value[index] = npy.complex(params[real_string], params[imag_string])
        
        return value[0] * value[1].conjugate()

    def Phasediff(self, params, wave1, wave2):
        value = npy.empty(2, "c16")
        
        # value[0] == wave1; value[1] == wave2
        for index, i in enumerate([wave1, wave2]):
            real_string = f"r.{self.__elm['e'][i]}.{self.__elm['l'][i]}.{self.__elm['m'][i]}"
            imag_string = f"i.{self.__elm['e'][i]}.{self.__elm['l'][i]}.{self.__elm['m'][i]}"
            value[index] = npy.complex(params[real_string], params[imag_string])
        phasediff = npy.arctan(npy.imag(value[0]*value[1].conjugate())/npy.real(value[0]*value[1].conjugate()))
        
        return phasediff