'''
Created on 10-May-2021

@author: bthakur

'''
import traceback
import inspect
import sympy
from sympy import (symbols, Symbol, latex, sympify, Eq, fraction,
                   srepr, solve
                   )

from ._latex_equation import LatexEquation
from xv.km.util import PintUnitSystem

from functools import reduce
import operator

class ProportionalEquation:
    def __init__(self, 
                 left_prop, 
                 right_params, 
                 const = 1, 
                 var_names = {}, 
                 **kwargs):
        
        self.proportional = Symbol(u'\U0000221d')
        self.implies = symbols(u'\U000021D2')
        
        if const == 1:
            const = 1
        self.const = const
        self.left_prop = left_prop
        self.right_params = right_params
        self.var_names = var_names
        self.desc = kwargs.get('desc', {})
        
        self.lhs = self._init_lhs()
        self.rhs = self._init_rhs()
        self.expr = self._init_expr()
        self.eq = self._init_eq()
        self.latex_eq = self._init_latex_equation(**kwargs)
        pass

    def get_statements(self, var_values = {}, annotate = True, html = True):
        statements = []
        if html:
            linebreak = "<br/>"
        else:
            linebreak = "\n"

        statements.append(self.get_var_lables(var_values = var_values))
            
        annotation = ""
        for item in self.right_params:
            if annotate:
                annotation = self._get_annotation(self.left_prop, item) + linebreak
            
            if html:
                statement = f"{annotation}${latex(self.left_prop)} {self.proportional} {latex(item)}$"
            else:
                statement = f"{annotation}{self.left_prop} {self.proportional} {item}"

            statements.append(statement)
            pass
        
        if html:
            statements.append(f"$\\implies {self.expr}$")
            
            const_annotation = f"the constant is ${latex(self.const)}$"
            
            statements.append(f"$\\implies {self.lhs} = {latex(self.const)} {self.rhs}$ \t\t {const_annotation}")
            if self.const == 1:
                statements.append(f"$\\implies {self.lhs} = {self.rhs}$")
        else:
            statements.append(f"\\implies {self.eq.lhs} {self.proportional} {self.eq.rhs/self.const}")
            const_annotation = f"the constant is {self.const}"
            statements.append(f"\\implies {self.eq} \t\t {const_annotation}")
            
        statements.append(f"$\\implies {latex(self.eq)}$")
        return f"{linebreak}{linebreak}".join(statements)
        pass    


    def get_solution_statements(self, 
                                var_values):
        
        left_symbol = self.get_left_symbol()

        eq_subs = self.subs_eq_quantities(var_values, [left_symbol])
        
        eq_subs_2 = self.subs_eq_quantities_2(var_values, [left_symbol])

        solved_val_with_units = self.solve_with_units(var_values, left_symbol)

        solution_statements = []

        stmt0 = f"$\\implies {latex(self.eq)}$"
        solution_statements.append(stmt0)

        stmt1 = f"$\\implies {latex(eq_subs)}$"
        if stmt1 != stmt0:
            solution_statements.append(stmt1)

        stmt2 = f"$\\implies {latex(eq_subs_2)}$"
        if stmt2 != stmt1:
            solution_statements.append(stmt2)
            
        solved_val_str = ", ".join([latex(v) for v in solved_val_with_units])
        stmt3 = f"$\\implies {latex(left_symbol)} = {solved_val_str}$"
        if stmt3 != stmt2:
            solution_statements.append(stmt3)
        
        pre_cal_soln_value = var_values.get(left_symbol, None)
        if not pre_cal_soln_value is None:
            stmt4 = f"$\\implies {latex(left_symbol)} = {latex(pre_cal_soln_value)}$"
            if stmt4 != stmt3:
                solution_statements.append(stmt3)

        return "\n\n".join(solution_statements)
        pass


    def get_left_symbol(self):
        left_symbols = self.left_prop.free_symbols
        return list(left_symbols)[0]

    def get_right_symbols(self):
        right_symbols = []
        for item in self.right_params:
            right_symbols.extend(item.free_symbols)
        return right_symbols
    
    def get_left_lable(self):
        left_vars = [self.get_left_symbol()]
        
        labels = [v for k, v in self.var_names.items()
                  if k in left_vars]
        
        return self._get_labels_str(labels)
        pass
    
    def get_right_lables(self, exclude_vars = []):
        include_vars = self.get_right_symbols()
            
        labels = [v for k, v in self.var_names.items()
                  if (k in include_vars and
                      k not in exclude_vars )]
        return self._get_labels_str(labels)
        pass

    def get_right_lables_with_values(self, var_values, exclude_vars = []):
        include_vars = [k for k in self.get_right_symbols() 
                        if not k in exclude_vars]
        
        result = []
        for k, label in self.var_names.items():
            if not k in include_vars:
                continue
            v = var_values.get(k, None)
            if not v is None:
                if (hasattr(v, 'magnitude') and 
                    hasattr(v, 'u')):
                    v_unit = str(v.u).replace("**", "^")
                    result.append(f"{label} = ${v.magnitude} \; {latex(v_unit)}$")
                else:
                    result.append(f"{label} = ${latex(v)}$")
            pass 
        
        return "\n\n".join(result)
        pass

    def get_symbols_labels_values(self, var_values, exclude_vars = []):
        result = []
        for k, label in self.var_names.items():
            if k in exclude_vars:
                continue
            
            v = var_values.get(k, None)
            if not v is None:
                if (hasattr(v, 'magnitude') and 
                    hasattr(v, 'u')):
                    v_unit = str(v.u).replace("**", "^")
                    result.append(f"{label} = ${v.magnitude} \; {latex(v_unit)}$")
                else:
                    result.append(f"{label} = ${latex(v)}$")
            pass 
        
        return "\n\n".join(result)
        pass

    def get_lables(self, exclude_vars = []):
        labels = [v for k, v in self.var_names.items()
                  if k not in exclude_vars]

        if len(labels) < 1:
            return ""
        elif len(labels) == 1:
            return labels[0]
        else:
            s = ", ".join(labels[:-1])
            s = f"{s} and {labels[-1]}"
            return s
        pass

    def subs_eq_quantities(self, var_quantities, exclude_vars = []):
        eq = self.eq.copy() 

        for k, v in var_quantities.items():
            if k in exclude_vars:
                continue
            if (hasattr(v, 'magnitude') and 
                hasattr(v, 'u')):
                v_unit = str(v.u).replace("**", "^")
                eq = eq.subs(k, Symbol(rf"\left( {v.magnitude} \; {v_unit} \right)"))
            else:
                eq = eq.subs(k, Symbol(str(v)))
        
        return eq
        pass
    
    def subs_eq_quantities_2(self, var_quantities, exclude_vars = []):
        eq = self.eq.copy() 

        for k, v in var_quantities.items():
            if k in exclude_vars:
                continue
            if (hasattr(v, 'magnitude') and 
                hasattr(v, 'u')):
                v_modified = v.magnitude * sympify(str(v.u))
            else:
                v_modified = v
                
            eq = eq.subs(k, v_modified)
        
        return eq
        pass
    
    def solve_eq_units(self, var_quantities, var_to_solve):
        eq = self.eq.copy() 

        for k, v in var_quantities.items():
            if k == var_to_solve:
                continue
            v_modified = sympify(str(v.u))
            eq = eq.subs(k, v_modified)
        
        solved_val = solve(eq, var_to_solve)
        return solved_val[0]
        pass
    
    def solve_with_units(self, var_quantities, var_to_solve):
        
        eq = self.subs_eq_quantities_2(var_quantities, [var_to_solve])
        solved_vals = solve(eq, var_to_solve)
        
        return solved_vals
        #return ",\n\n ".join([f"${latex(v)}$" for v in solved_vals])
        pass
    
    def solve(self, var_quantities, var_to_solve):
        
        eq = self.eq.copy() 
        var_values = {}

        for k, v in var_quantities.items():
            if k == var_to_solve:
                continue
            
            if hasattr(v, 'magnitude'):
                v = v.magnitude
            var_values[k] = v

        eq = eq.subs(var_values)    
        solved_val = solve(eq, var_to_solve)
            
        return ", ".join([f"{v}" for v in solved_val])
        pass
    
    
    def compute(self, var_values, var_to_solve):
        try:
            eq = self.eq.copy()
            for k, v in var_values.items():
                if k == var_to_solve:
                    continue
                #print(type(v).__name__)  #Quantity
                eq = eq.subs(k, v.magnitude * k)
                eq_s = srepr(eq).replace(srepr(k), srepr(v.u))
                eq = sympify(eq_s, evaluate = False)
                pass
            
            ps = PintUnitSystem()
            solved_val = solve(eq, var_to_solve)
            
            solved_val = ps.get_quantity(str(solved_val[0]))
            
            solved_val = solved_val.to_base_units()
            return eq, solved_val
            pass
        except Exception as e:
            try:
                sError = f"""{inspect.stack()[0][3]} 
                First level Error: {e}
                eq: {self.eq}
                
                var_values: {var_values}
                
                var_to_solve: {var_to_solve}
                """
                if self.verbose:
                    print(sError)
                    print(traceback.format_exc())

                eq = self.eq.subs(var_values)
                solved_val = solve(eq, var_to_solve)[0]
                return eq, solved_val
                pass
            except Exception as e:
                sError = f"""{inspect.stack()[0][3]} Error: {e}
                eq: {self.eq}
                """
                if self.verbose:
                    print(sError)
                    print(traceback.format_exc())
        
        return eq, solved_val
        pass
    
    def transform(self, new_left_symbol):
        if new_left_symbol in self.left_prop.free_symbols:
            return self
        
        new_left_prop = [item for item in self.right_params 
                        if new_left_symbol in item.free_symbols]
        new_left_prop = 1 / new_left_prop[0]
        
        new_right_params = [item for item in self.right_params 
                        if not new_left_symbol in item.free_symbols]
        
        new_right_params.insert(0, 1 / self.left_prop)
        
        const = self.const
        
        if self._is_inverse(new_left_prop, new_left_symbol):
            new_left_prop = 1 /new_left_prop
            new_right_params = [1/item for item in new_right_params]
            const =  1 / const
        
        npe = ProportionalEquation(new_left_prop,
                                 new_right_params, 
                                 const, 
                                 var_names = self.var_names
                                 )
        return npe
        pass
    
    def _get_annotation(self, lterm, rterm):
        lterm_symb = list(lterm.free_symbols)[0]
        rterm_symb = list(rterm.free_symbols)[0]
        lterm_name = self.var_names.get(lterm_symb, lterm_symb).lower()
        rterm_name = self.var_names.get(rterm_symb, rterm_symb).lower()

        lterm_inverse = self._is_inverse(lterm, lterm_symb)
        rterm_inverse = self._is_inverse(rterm, rterm_symb)

        if (lterm_inverse == rterm_inverse):
            annotation = f"More {rterm_name} means more {lterm_name}"
        else:
            annotation = f"Less {rterm_name} means more {lterm_name}"
            pass

        return annotation
        pass
    
    def _is_inverse(self, term, symb):
        num_term, _ = fraction(term)
        if not symb in num_term.free_symbols:
            return True
        return False
    
    def get_desc(self, html = True):
        if len(self.desc) < 1:
            return ""
        
        if html:
            return self._get_dict_html(self.desc)
        
        return self._get_dict_str(self.desc)
        pass
    
    def get_var_lables(self, var_values = {}, html = True):
        if len(self.var_names) < 1:
            return ""
        
        if html:
            return self._get_dict_html(self.var_names,
                                       var_values = var_values)
        
        return self._get_dict_str(self.var_names, 
                                  var_values = var_values)
        pass
    
    def _init_lhs(self):
        lhs = f"{latex(self.left_prop)}"
        return lhs
        pass
    
    def _init_rhs(self):
        rhs = [f"{latex(p)}" for p in self.right_params]
        #rhs.insert(0, f"{latex(self.const)}")
        rhs = " ".join(rhs)
        return rhs
        pass

    def _init_expr(self):
        s = f"{self.lhs} {self.proportional} {self.rhs}"
        return s.strip()
    
    
    def _init_latex_equation(self, **kwargs):
        to_latex = kwargs.get('to_latex', False)
        if not to_latex:
            latex_eq = LatexEquation.from_eq(
                            self.eq, var_names = self.var_names)
            return latex_eq
        
        kwargs['to_latex'] = False
        
        latex_eq = LatexEquation(self.lhs, 
                                 self._get_rhs_with_const(), 
                                 eq = self.eq, 
                                 **kwargs)
        return latex_eq
        pass
    
    def _get_rhs_with_const(self):
        if self.const == 1:
            return self.rhs
        return f"{self.const}{self.rhs}"
        pass
    
    def _init_eq(self):
        rhs_eq = [p for p in self.right_params]
        if self.const != 1:
            rhs_eq.insert(0, self.const)
        
        #rhs_eq = math.prod(rhs_eq)
        rhs_eq = reduce(operator.mul, rhs_eq, 1)

        
        eq = Eq(self.left_prop, rhs_eq)
        return eq
        pass
    

    def _get_dict_html(self, dct, var_values = {}):
        s = []
        for k, v in dct.items():
            k_value = var_values.get(k, None)
            if not k_value is None:
                s.append(f"${latex(k)}$ = {v} = ${latex(k_value)}$ ")
            else:
                s.append(f"${latex(k)}$ = {v}")
            pass
        return "<br/>".join(s).strip()
        pass
    
    def _get_dict_str(self, dct, var_values = {}):
        s = []
        for k, v in dct.items():
            k_value = var_values.get(k, None)
            if not k_value is None:
                s.append(f"{k} = {v} = {k_value}")
            else:
                s.append(f"{k} = {v}")
            pass
        return "\n".join(s).strip()
        pass
    
    def _get_labels_str(self, labels):
        if len(labels) < 1:
            return ""
        elif len(labels) == 1:
            return labels[0]
        else:
            s = ", ".join(labels[:-1])
            s = f"{s} and {labels[-1]}"
            return s
        pass

    def _repr_html_(self):
        line_break = "<br/>"
        s = f"""
        ${self.lhs} {self.proportional} {self.rhs}$
        {line_break}
        {line_break}
        $\\implies {self.lhs} = {self._get_rhs_with_const()}$
        """
        return s.replace("\n", "").strip()
        pass

    def _repr_latex_(self):
        s = f"""
        ${self.lhs} {self.proportional} {self.rhs}$
        $\\implies {self.lhs} = {self._get_rhs_with_const()}$
        """
        return s.strip()
        pass

    def __str__(self):
        s = f"""
        {self.lhs} {self.proportional} {self.rhs}
        \\implies {latex(self.eq)}
        """
        return s.strip()
        pass
    
    pass

#class ProportionalEquation:
    