'''
Created on 23-May-2021

@author: bthakur
'''
from sympy import Symbol, Eq, sympify, latex
import random
from xv.km.util import latex_fraction
from xv.km.kelements import KPrintElement

from ._physical_quantity import PhysicalQuantity
from sympy.physics import units

class PhysicalQuantitySystem(object):
    '''
    defines a set of variables that related by a rule of proportion
    '''

    class _SystemState:
        def __init__(self, name, member_scale_factors, abbrev = None, prev_state = None):
            '''
            member_scale_factors is dict of members and their values
            '''
            self.name = name
            if not abbrev is None:
                self.abbrev = abbrev
            else:
                self.abbrev = self.name
                
            self.prev_state = prev_state
            self.member_scale_factors = member_scale_factors
            pass
        
        def reset_member_state(self, member, scale_factor):
            self.member_scale_factors[member] = scale_factor
            pass

        def __repr__(self):
            return str(self)
        
        def __str__(self):
            s = type(self).__name__
            s = s + "\t" + f"{self.abbrev}: {self.name}" 
            s = s + "\n" + f"prev state: {self.prev_state}" 
            
            s = s + self._str_member_scale_factor() 
            return s
            pass
        
        def  _str_member_scale_factor(self):
            s = ""
            for member, scale_factor in self.member_scale_factors.items():
                s = s + "\n" + f"{member}: {scale_factor}"
            return s
            pass
    
        pass
    #class _SystemState

    class _SystemEquation:
        def __init__(self, name, eq, abbrev = None):
            '''
            '''
            self.name = name
            if not abbrev is None:
                self.abbrev = abbrev
            else:
                self.abbrev = self.name
                
            self.eq = eq
            pass
        
        def __repr__(self):
            s = f"{self.abbrev}: {self.name}: ${self.eq.__repr__()}$" 
            return s
        
        def _repr_latex_(self):
            s = f"{self.abbrev}: {self.name}: ${self.eq._repr_latex_()}$" 
            return s
        
        
        def __str__(self):
            s = f"{self.abbrev}: {self.name}: {self.eq}" 
            return s
            pass
    
        pass
    #class _SystemState

    class _SystemProportionality:
        def __init__(self, left_term, *right_terms, constant = 1):
            self.left_term = left_term
            self.right_terms = right_terms
            self.constant = constant
            pass
 
        def getDefaultLabelNameList(self):
            terms_list = [(self.left_term.name, self.left_term.abbrev)]
            for right_term in self.right_terms:
                if (isinstance(right_term, tuple) or 
                    isinstance(right_term, list)):
                    terms_list.append((right_term[0].name, right_term[0].abbrev))
                else:
                    terms_list.append((right_term.name, right_term.abbrev))
            
            return terms_list
            pass

        def getDefaultEq(self):
            right_terms_prod = 1
            for right_term in self.right_terms:
                if (isinstance(right_term, tuple) or 
                    isinstance(right_term, list)):
                    rtb, rte = right_term
                    if rte < 0:
                        right_terms_prod = right_terms_prod * (1 / rtb.abbrev ** (0-rte))
                    else:
                        right_terms_prod = right_terms_prod * rtb.abbrev ** rte
                else:
                    right_terms_prod = right_terms_prod * right_term.abbrev
            
            if not self._is_constant_one():
                if hasattr(self.constant, 'abbrev'):
                    right_terms_prod = self.constant.abbrev * right_terms_prod
                else:
                    right_terms_prod = self.constant * right_terms_prod
                pass
            eq = Eq(self.left_term, right_terms_prod)
            return eq
            pass       

        def _is_constant_one(self):
            if (self.constant.dimension == units.dimensions.Dimension(1) and 
                self.constant.scale_factor == 1):
                return True
            return False
            pass

        def get_expression(self):
            proportional = Symbol(u'\U0000221d')
            implies = Symbol(u'\U000021D2')
    
            left_term = self.left_term
            right_terms = self.right_terms
            
            eq_latex = KPrintElement.getLatex(self.getDefaultEq())
            expr = ""
            
            right_terms_prod = ""
            for right_term in right_terms:
                if (isinstance(right_term, tuple) or 
                    isinstance(right_term, list)):
                    rtb, rte = right_term
                    if rte < 0:
                        rt_expr = latex_fraction("1", rtb.abbrev ** (0-rte))
                        pass
                    else:
                        rt_expr = rtb.abbrev ** rte
                        pass
                    
                    expr = f"""{expr}
                    ${left_term.abbrev} {proportional} {rt_expr}$
                    """
                    right_terms_prod = f"""{right_terms_prod} * {rt_expr}"""
                else:
                    expr = f"""{expr}
                    ${left_term.abbrev} {proportional} {right_term.abbrev}$
                    """
                    right_terms_prod = f"""{right_terms_prod} * {right_term.abbrev}"""
                   
            expr = f"""{expr}
            ${implies} {left_term.abbrev} {proportional}{right_terms_prod}$
            
            ${implies} {eq_latex}$
            """
            return expr.replace("**", "^").replace("*", "")
            pass
   
        pass
    #class _SystemProportionality:

    def __init__(self, name, abbrev = None, **kwargs):
        '''
        params:
        '''
        self.name = name
        if not abbrev is None:
            self.abbrev = abbrev
        else:
            self.abbrev = self.name
        
        self._member_names = []
        self._member_abbrev = []
        
        self._system_states = {}
        self._system_equations = {}
        
        self._proportionality = None
        self._statement = None
        
        self.verbose = kwargs.get('verbose', True)
        pass

    def add_statement(self, statement):
        self._statement = statement
        pass

    def get_statement(self):
        if self._statement is None:
            self._init_default_statement()
        return self._statement
        pass

    def get_proportionality_expression(self):
        return self._proportionality.get_expression()
        pass

    def _init_default_statement_NOT_WORKING(self):
        try:
            if not self._proportionality is None:
                left_term = self._proportionality.get('left_term')
                name_label_list = [(left_term.name, left_term.abbrev)]
                
                name_label_list.extend(
                    [ (right_term.name, right_term.abbrev)
                        for right_term in self._proportionality.get('right_terms') ]
                    )
                pass
            else:
                name_label_list = [
                    (self._member_names[i], self._member_abbrev[i])
                    for i in range(len(self._member_names))
                    ]
                pass

            random.shuffle(name_label_list)
            
            stmt_str = ", ".join ([
                f"{label} ${name}$"
                for (label, name) in name_label_list[:-1]
                ])
            
            stmt_str = f"""{stmt_str} and {name_label_list[-1][0]} ${name_label_list[-1][1]}$"""
            
            statement = f"""
            {self.name}:
            Write the expression for relation among {stmt_str}
            """

            self._statement = statement
        except Exception as e:
            print(f"Error _init_default_statement: {e}")
            pass
        pass

    def _init_default_statement(self):
        try:
            name_label_list_NOT_USED = [
                (self._member_names[i], self._member_abbrev[i])
                for i in range(len(self._member_names))
                ]
            
            name_label_list = self._proportionality.getDefaultLabelNameList()
            
            random.shuffle(name_label_list)
            
            stmt_str = ", ".join ([
                f"{label} ${name}$"
                for (label, name) in name_label_list[:-1]
                ])
            
            stmt_str = f"""{stmt_str} and {name_label_list[-1][0]} ${name_label_list[-1][1]}$"""
            
            statement = f"""
            {self.name}:
            Write the expression for relation among {stmt_str}
            """

            self._statement = statement
        except Exception as e:
            print(f"Error _init_default_statement: {e}")
            pass
        pass

    def set_proportionality(self, left_term, *right_terms, constant = 1):
        self._proportionality = PhysicalQuantitySystem._SystemProportionality(
            left_term, *right_terms, constant = constant
            )
            
        self.add_system_equation('default', 
                    self._proportionality.getDefaultEq()
                                 )
        pass

    def get_proportionality(self):
        return self._proportionality 
        pass

    def add_physical_quantity(self, physical_quantity, overwrite = False):
        if not overwrite:
            if self._is_member(physical_quantity):
                raise Exception(f"""
                The physical_quantity {physical_quantity} already exists.
                Use 'overwrite = True' if you want to add again
                """)
        
        setattr(self, str(physical_quantity.name), physical_quantity)
        if hasattr(physical_quantity, 'abbrev'):
            setattr(self, str(physical_quantity.abbrev), physical_quantity)
            
        self._add_members(physical_quantity)
        pass

    def get_system_state(self, name):
        for system_state in self._system_states.values():
            if (system_state.name == name or 
                system_state.abbrev == name):
                return system_state
        return None
        pass

    
    def add_system_state(self, name, member_scale_factors, 
                         abbrev = None, 
                         prev_state = None, 
                         overwrite = False):
        if abbrev is None:
            abbrev = name
            
        if self._is_system_state(abbrev, name):
            if not overwrite:
                raise Exception(f"""
                The {name} or {abbrev} already exists in system states.
                Use 'overwrite = True' if you want to add again
                """)
            else:
                self.remove_system_state(abbrev, name)
               
        if not prev_state is None:
            if not self._is_system_state(prev_state):
                raise Exception(f"""
                The prev state '{prev_state}' does not exist in system states.
                """)
            
        self._system_states[abbrev] = PhysicalQuantitySystem._SystemState(
                                        name, 
                                        member_scale_factors,
                                        abbrev = abbrev, 
                                        prev_state = prev_state
                                    )
        pass

    def get_system_equation(self, name = 'default'):
        for system_equation in self._system_equations.values():
            if (system_equation.name == name or 
                system_equation.abbrev == name):
                return system_equation.eq
        return None
        pass
    
    def add_system_equation(self, name, eq, 
                         abbrev = None, 
                         prev_state = None, 
                         overwrite = False):
        if abbrev is None:
            abbrev = name
            
        if self._is_system_equation(abbrev, name):
            if not overwrite:
                raise Exception(f"""
                The {name} or {abbrev} already exists in system states.
                Use 'overwrite = True' if you want to add again
                """)
            else:
                self.remove_system_equation(abbrev, name)
                
        self._system_equations[name] = PhysicalQuantitySystem._SystemEquation(
                                        name, 
                                        eq,
                                        abbrev = abbrev
                                    )
        pass

    def remove_system_state(self, abbrev, name = None):
        is_deleted = False
        if abbrev in self._system_states.keys():
            del self._system_states[abbrev]
            is_deleted = True

        if not name is None:
            if name in self._system_states.keys():
                del self._system_states[name]
                is_deleted = True
        
        if not is_deleted:  
            if not name is None:
                print(f"the state '{abbrev}' does not exist")
            else:
                print(f"the state '{abbrev}' or '{name}' does not exist")
        pass

    
    def remove_system_equation(self, abbrev, name = None):
        is_deleted = False
        if abbrev in self._system_equations.keys():
            del self._system_equations[abbrev]
            is_deleted = True

        if not name is None:
            if name in self._system_equations.keys():
                del self._system_equations[name]
                is_deleted = True
        
        if not is_deleted:  
            if not name is None:
                print(f"the equation '{abbrev}' does not exist")
            else:
                print(f"the equation '{abbrev}' or '{name}' does not exist")
        pass

    def _add_members(self, physical_quantity):
        self._member_names.append(str(physical_quantity.name))
        self._member_abbrev.append(str(self._get_member_abbrev(physical_quantity)))
    
    def _is_member(self, physical_quantity):
        if str(physical_quantity.name) in self._member_names:
            return True
        
        if str(self._get_member_abbrev(physical_quantity)) in self._member_abbrev:
            return True
        
        return False
        pass

    def _is_system_state(self, abbrev, name = None):
        if abbrev in self._system_states.keys():
            return True
        
        if not name is None:
            if name in self._system_states.keys():
                return True
        
        return False
        pass

    def _is_system_equation(self, abbrev, name = None):
        if abbrev in self._system_equations.keys():
            return True
        
        if not name is None:
            if name in self._system_equations.keys():
                return True
        
        return False
        pass

    def _get_member_abbrev(self, physical_quantity):
        if hasattr(physical_quantity, 'abbrev'):
            return physical_quantity.abbrev
        return physical_quantity.name
        pass

    def __repr__(self):
        s = type(self).__name__
        s = s + "\n\t" + "name" + ": " + str(self.name) + f"({self.abbrev})" 
        s = s + "\n\t" + "members" + ": " + str(self._member_abbrev) 
        s = s + "\n\t" + f"states: {list(self._system_states.keys())}"
        s = s + "\n\t" + f"equations: {list(self._system_equations.keys())}"
        s = s + "\n\t" + f"proportionality: {self.get_proportionality()}"
        return s
    
    def __str__(self):
        s = type(self).__name__
        s = s + "\n\t" + "name" + ": " + str(self.name) + f"({self.abbrev})" 
        s = s + "\n\t" + "members" + ": " + str(self._member_abbrev) 
        s = s + "\n\t" + f"states: {list(self._system_states.keys())}"
        s = s + "\n\t" + f"equations: {list(self._system_equations.keys())}"
        s = s + "\n\t" + f"proportionality: {self.get_proportionality()}"
        return s
        pass
    
    pass
#class PhysicalQuantitySystem(