from sympy import Symbol, Rational, Eq, solve, latex
import numpy as np

class QuantumParticle:

    def __init__(self, **kwargs):
        self.verbose = kwargs.get('verbose', False)
        # self._prefix = kwargs.get('prefix', None)
        
        self._m = self.init_symbol('m', **kwargs)
        self._p = self.init_symbol('p', **kwargs)
        self._v = self.init_symbol('v', **kwargs)
        self._ke = self.init_symbol('KE', **kwargs)
        
        self._m_val = kwargs.get('m', self._m)
        self._p_val = kwargs.get('p', self._p)
        self._v_val = self._init_v_val(**kwargs)
        self._ke_val = self._init_ke_val(**kwargs)
        pass
        
    def _init_v_val(self, **kwargs):
        v_input = kwargs.get('v', None)
        if isinstance(v_input, list):
            v_input = np.array(v_input)

        mass_val = self.get_mass_value()
        momentum_val = self.get_momentum_value()
        if isinstance(momentum_val, list):
            momentum_val = np.array(momentum_val)

        v_computed = momentum_val / mass_val

        if v_input is None:
            return v_computed

        if (self.get_mass() == mass_val or 
            self.get_velocity() == momentum_val):
            return v_input
        
        if list(v_input) != list(v_computed):
            sError = f'v_input = {v_input} and v_computed = {v_computed} do not match.'
            raise Exception(sError) 
        
        return v_input
        
        pass

    def _init_ke_val(self, **kwargs):
        ke_input = kwargs.get('KE', None)
        mass_val = self.get_mass_value()
        velocity_val = self.get_velocity_value()
        if isinstance(velocity_val, list):
            velocity_val = np.array(velocity_val)
        
        if isinstance(velocity_val, np.ndarray):
            ke_computed = Rational(1, 2)* sum(velocity_val)**2 / mass_val
        else:
            ke_computed = Rational(1, 2) * velocity_val**2 / mass_val

        if ke_input is None:
            return ke_computed

        if (self.get_mass() == mass_val or 
            self.get_velocity() == velocity_val):
            return ke_input
        
        if ke_input != ke_computed:
            sError = f'ke_input = {ke_input} and ke_computed = {ke_computed} do not match.'
            raise Exception(sError) 

        return ke_input

        pass
    
    def get_value_dict(self):
        value_dict = {
            self.get_mass(): self.get_mass_value(),
            self.get_momentum(): self.get_momentum_value(),
            self.get_velocity(): self.get_velocity_value(),

        }
        return value_dict
    
    def get_mass(self):
        return self._m
    
    def get_mass_value(self):
        return self._m_val

    def get_momentum(self):
        return self._p
    
    def get_momentum_value(self):
        return self._p_val

    def get_velocity(self):
        return self._v
    
    def get_velocity_value(self):
        return self._v_val

    def get_kinetic_energy(self):
        return self._ke
    
    def get_kinetic_energy_value(self):
        return self._ke_val

    def compute_velocity(self):
        return self.get_momentum() / self.get_mass()
    
    def init_symbol(self, s, **kwargs):
        prefix = kwargs.get(f'prefix_{s}', None)
        if prefix is None:
            prefix = kwargs.get('prefix', None)
            
        if prefix is None:
            return Symbol(s)
        
        return Symbol(f'{s}_{{{prefix}}}')
        pass

    def __str__(self):
        cls = self.__class__
        s = []
        s.append(f"{id(self)}@{cls.__name__}")
        s.append(f"mass: {self.get_mass()} = {self.get_mass_value()}")
        s.append(f"momentum: {self.get_momentum()} = {self.get_momentum_value()}")
        s.append(f"velocity: {self.get_velocity()} = {self.get_velocity_value()}")
        s.append(f"kinetic_energy: ${latex(self.get_kinetic_energy())} = {latex(self.get_kinetic_energy_value())}$")
        s = "\n".join(s)
        return s

    def _repr_html_(self):
        cls = self.__class__
        s = []
        s.append(f"{id(self)}@{cls.__name__}")
        s.append(f"mass: ${latex(self.get_mass())} = {latex(self.get_mass_value())}$")
        s.append(f"momentum: ${latex(self.get_momentum())} = {latex(self.get_momentum_value())}$")
        s.append(f"velocity: ${latex(self.get_velocity())} = {latex(self.get_velocity_value())}$")
        s.append(f"kinetic_energy: ${latex(self.get_kinetic_energy())} = {latex(self.get_kinetic_energy_value())}$")
        s = "<br/>".join(s)
        return s    
