'''
Created on 12-Jul-2021

@author: rishi_priyansh
'''
import inspect
import traceback
import random

from sympy import (latex, Symbol, symbols, UnevaluatedExpr as ueval,
                   Eq, Matrix,
                   )
import sympy as sp
from xv.km.util import SympyUnitSystem

# _r_circuit_helper
# RCircuitHelper

from ._base_circuit_helper import BaseCircuitHelper
from mkl_random.mklrand import rand

class RCircuitHelper(BaseCircuitHelper):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        pass
    
    def get_problem_circuit_R_resultant(self, **kwargs):
        
        sInfo = f"""Processing {self.__class__.__name__}.{inspect.stack()[0][3]}: """
        self._log(sInfo.strip())

        index = kwargs.get('index', None)
        
        try:
            
            if index is None:
                random_choices = list(range(1, 2))
        
                identifier = f"{self.__class__.__name__}{id(self)}.{inspect.stack()[0][3]}.random_choices"
                random_choice = self.get_no_repeat_random_choice(identifier, random_choices)
            else:
                random_choice = index
            
            if random_choice == 1:
                return self._get_mixed_R_circuit(**kwargs)
    
            return self._get_mixed_R_circuit(**kwargs)
        
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            """
            self._log(sError)
            self._log(traceback.format_exc())
            pass
        
        return None

    def _get_mixed_R_circuit(self, **kwargs):
        import schemdraw
        import matplotlib
        matplotlib.use('Agg') 

        R = Symbol('R')
        R_unit = self._ss.to_unit('ohm') 
        
        arrange_types = {'s':'series', 'p':'parallel'}
        
        R_sets = []
        d = schemdraw.Drawing()
        
        prev_arrange_type = None
        
        for idx in range(random.randint(2, 4)):
            if prev_arrange_type == 's':
                arrange_type = "p"
            else:
                arrange_type = random.choice(list(arrange_types.keys()))
                
            if prev_arrange_type == 's':
                min_count = 1
            else:
                min_count = 2
             
            R_set = []
            for R_idx in range(random.randint(min_count, 4)):
                R_val = ueval(random.randint(1, 5)) * R_unit
                R_set.append(R_val)
                pass
            
            R_sets.append([R_set, arrange_type])
            
            d_set = self._draw_R_set(d, R_set, arrange_type)
            prev_arrange_type = arrange_type
            pass
        
        svg_str = self._convert_d_to_svg_str(d)

        prob_s = f"""
        Draw the simplified circuit and find voltage, resistance and current in the circuit.
        {svg_str}
        """

        soln_stmts = []
        
        
        R_sets_ltx = []
        R_sums_ueval = []
        R_sums = []
        for R_set, arrange_type in R_sets:
            arrange_desc = arrange_types.get(arrange_type)
            R_items_ltx = [f"{latex(item)}" for item in R_set]
            R_set_ltx = ", ".join(R_items_ltx)
            R_set_ltx = f" \\left( {R_set_ltx} \\text{{ in {arrange_desc}}} \\right)"
            R_sets_ltx.append(R_set_ltx)
            
            if arrange_type == "s":
                R_sum_ueval = " + ".join(R_items_ltx)
                R_sums_ueval.append(R_sum_ueval)

                R_sums.append(sp.Add(*R_set).simplify())
                
            else:
                R_items_inv_ltx = [f"\\frac{{1}}{{{latex(item)}}}" for item in R_set]
                R_sum_ueval = " + ".join(R_items_inv_ltx)
                R_sum_ueval = f"\\frac{{ 1 }} {{ {R_sum_ueval} }} "
                R_sums_ueval.append(R_sum_ueval)
            
                R_set_inv_sum = sp.Add(*[1/item for item in R_set])
                R_sums.append(1 / R_set_inv_sum.simplify())
                pass
            
            pass
        
        R_sets_ltx = " \\\\ \\\\ & + ".join(R_sets_ltx)
        R_sums_ueval = " + ".join(R_sums_ueval)
        
        R_sums_ltx = " + ".join([f"{latex(item)}" for item in R_sums])
        
        R_val = sum(R_sums).simplify()

        d_simplified = schemdraw.Drawing()
        self._draw_series_R_set(d_simplified, R_sums)
        svg_str_simplified = self._convert_d_to_svg_str(d_simplified)
        

        soln_stmt = f"""
        We have the resistance sets in the series.
        
        $\\displaystyle 
        \\begin{{aligned}}
        
        {latex(R)} 
        = & \\; {R_sets_ltx}
        \\\\ \\\\
      
        = & \\; {R_sums_ueval}
        \\\\ \\\\
      
        = & \\; {R_sums_ltx}
        \\\\ \\\\
        
        = & \\; {latex(R_val)}
        
        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())
        

        soln_s = "\n\n".join(soln_stmts)
        
        answer_s = f"""
        $\\displaystyle 
        {latex(R_val)}
        $
        {svg_str_simplified}
        """
        
        result = {
            'prob_s': prob_s.strip(),
            'soln_s': soln_s.strip(),
            'answer_s': answer_s.strip(),
            }
        return result


    def _draw_R_set(self, d, R_set, arrange_type):
        if arrange_type == 's':
            return self._draw_series_R_set(d, R_set)
        
        return self._draw_parallel_R_set(d, R_set)
        pass

    def _draw_series_R_set(self, d, R_set):
        import schemdraw.elements as elm
        
        for R_temp in R_set:
            d.add(elm.Resistor().label(f'${latex(R_temp)}$', loc = "top"))
        
        return d
        pass

    def _draw_parallel_R_set(self, d, R_set):
        import schemdraw.elements as elm
        
        std_length = d.unit * .75
        
        start_line = d.add( elm.Line(d='right', l = std_length ))
        
        for R_idx, R_temp in enumerate(R_set):
            d.push()
            L_length = std_length * (R_idx + 1)
            
            L_up = d.add(elm.Line(d='up', l=L_length))
            d.add(elm.Resistor(d='right', start = L_up.end
                               ).label(f'${latex(R_temp)}$', 
                                       loc = "top",
                                       ))
            L_up = d.add(elm.Line(d='down', l=L_length))
            if R_idx < len(R_set) - 1:
                d.pop()
            pass
        
        end_line = d.add( elm.Line(d='right', l = std_length ))
        return d
        pass


    
    def get_problem_circuit_R_bridge(self, **kwargs):
        
        sInfo = f"""Processing {self.__class__.__name__}.{inspect.stack()[0][3]}: """
        self._log(sInfo.strip())

        index = kwargs.get('index', None)
        
        try:
            
            if index is None:
                random_choices = list(range(1, 2))
        
                identifier = f"{self.__class__.__name__}{id(self)}.{inspect.stack()[0][3]}.random_choices"
                random_choice = self.get_no_repeat_random_choice(identifier, random_choices)
            else:
                random_choice = index
            
            if random_choice == 1:
                return self._get_bridge_R_circuit(**kwargs)
    
            return self._get_bridge_R_circuit(**kwargs)
        
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}: {e}
            """
            self._log(sError)
            self._log(traceback.format_exc())
            pass
        
        return None

    
    def _get_bridge_R_circuit(self, **kwargs):
        E0 = Symbol('E_0')
        V = Symbol('V')
        I, I0, I1, I2, I3 = symbols("I, I_0, I_1, I_2, I_3")
        R, R1, R2, R3, R4, R5 = symbols("R, R_1, R_2, R_3, R_4, R_5")
        R_unit = self._ss.to_unit('ohm') 
        I_unit = self._ss.to_unit('ampere') 
        V_unit = self._ss.to_unit('volt') 
        
        
        
        val_choice = random.randint(0, 6)
        
        if val_choice == 0:
            val_dict = {item: item for item in [R1, R2, R3, R4, R5]}
            val_dict[E0] = E0
        elif val_choice == 1:
            r1_val = random.randint(1, 5) * R_unit 
            r2_val = random.randint(1, 5) * R_unit 
            
            val_dict = {
                R1: r1_val, 
                R2: r1_val, 
                R3: r2_val, 
                R4: r2_val, 
                R5: random.choice([r1_val, r2_val, r1_val + r2_val, abs(r1_val - r2_val)])
                }
            val_dict[E0] = random.randint(4, 24) * V_unit
        elif val_choice == 2:
            r1_val = random.randint(1, 5) * R_unit 
            r2_val = random.randint(1, 5) * R_unit 
            
            val_dict = {
                R1: r1_val, 
                R2: r2_val, 
                R3: r1_val, 
                R4: r2_val, 
                R5: random.choice([r1_val, r2_val, r1_val + r2_val, abs(r1_val - r2_val)])
                }
            val_dict[E0] = random.randint(4, 24) * V_unit
        else:
            val_dict = {item: random.randint(1, 5) * R_unit 
                                  for item in [R1, R2, R3, R4, R5]}
            val_dict[E0] = random.randint(4, 24) * V_unit
            pass
        
        svg_str = self._draw_bridge(val_dict, E0,
                                    R1, R2, R3, R4, R5,  
                                    I0, I1, I3,  
                                    show_current = False,
                                    **kwargs)

        prob_s = f"""
        val_choice = {val_choice}
        {svg_str}
        Find the equivalent resistance of the bridge betweeen $B$ and $E$.
        Also find the current ${latex(I0)}$.
        """

        soln_stmts = []
        
        
        svg_str_with_current = self._draw_bridge(val_dict, E0,
                                                 R1, R2, R3, R4, R5, 
                                                 I0, I1, I3,  
                                                 show_current = True,
                                                 **kwargs)
        

        soln_stmt = f"""
        Let us show the current flowing and equivalent voltage $({latex(I*R)})
        $ for each branch of the circuit.
        
        {svg_str_with_current}
        """
        soln_stmts.append(soln_stmt.strip())
        

        soln_stmt = f"""
        In circuit $BCDB$,
        
        $\\displaystyle 
        \\begin{{aligned}}
        
        & {latex(I1)}{latex(val_dict[R1])} 
        + \\left( {latex(I1)} - {latex(I3)} \\right) {latex(val_dict[R5])} 
        - \\left( {latex(I0)} - {latex(I1)} \\right) {latex(val_dict[R2])} 
        =  0
        \\\\ \\\\
        
        \\implies
        & - {latex(val_dict[R2])} \\cdot {latex(I0)}
        + \\left( {latex(val_dict[R1] + val_dict[R2] + val_dict[R5])} \\right) \\cdot {latex(I1)}
        - {latex(val_dict[R5])} \\cdot {latex(I3)}
        =  0
        \\qquad \\cdots \\text{{ (i) }}
        
        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())
        

        soln_stmt = f"""
        In circuit $CEDC$,
        
        $\\displaystyle 
        \\begin{{aligned}}
        
        & {latex(I3)}{latex(val_dict[R3])} 
        - \\left( {latex(I0)} - {latex(I3)} \\right) {latex(val_dict[R4])} 
        - \\left( {latex(I1)} - {latex(I3)} \\right) {latex(val_dict[R5])} 
        =  0
        \\\\ \\\\
        
        \\implies
        & - {latex(val_dict[R4])} \\cdot {latex(I0)}
        - {latex(val_dict[R5])} \\cdot {latex(I1)}
        + \\left( {latex(val_dict[R3] + val_dict[R4] + val_dict[R5])} \\right) \\cdot {latex(I3)}
        =  0
        \\qquad \\cdots \\text{{ (ii) }}
        
        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())
        

        soln_stmt = f"""
        In circuit $ABCEFGKA$,
        
        $\\displaystyle 
        \\begin{{aligned}}
        
        & {latex(I1)}{latex(val_dict[R1])} 
        + {latex(I3)}{latex(val_dict[R3])}
        - {latex(val_dict[E0])} 
        =  0
        \\\\ \\\\
        
        \\implies
        & {latex(val_dict[R1])} \\cdot {latex(I1)}
        + {latex(val_dict[R3])} \\cdot {latex(I3)}
        = {latex(val_dict[E0])} 
        \\qquad \\cdots \\text{{ (iii) }}
        
        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())

        a1, b1, c1, a2, b2, c2, a3, b3, c3 = symbols("a1, b1, c1, a2, b2, c2, a3, b3, c3") 
        
        eq_BCDB = Eq( a1 * I0 + b1 * I1 + c1 * I3 , 0)
        eq_CEDC = Eq( a2 * I0 + b2 * I1 + c2 * I3 , 0)
        eq_ABCEFGKA = Eq(a3 * I0 + b3 * I1 + c3 * I3 , val_dict[E0])
        
        subs_dict = {
            a1: - val_dict[R2],
            b1: val_dict[R1] + val_dict[R2] + val_dict[R5],
            c1: val_dict[R5],
            a2: - val_dict[R4],
            b2: - val_dict[R4],
            c2: val_dict[R3] + val_dict[R4] + val_dict[R5],
            a3: 0,
            b3: val_dict[R1],
            c3: val_dict[R3],
            }
        
        eqs = [eq_BCDB, eq_CEDC, eq_ABCEFGKA]
        eq_vars = [I0, I1, I3]
        
        eqs_mod = [item.subs(subs_dict) for item in eqs]
        I_mod_solns = sp.solve(eqs_mod, *eq_vars, dict = True)[0]
        
        I0_mod_val = I_mod_solns.get(I0, None)
        
        I0_val = self._ss.convert_to(I0_mod_val, [I_unit])
        
        R_inv_val = (abs(I0_mod_val) / val_dict[E0]).simplify() 
        
        R_val = (1 / R_inv_val).simplify()
        
        soln_stmt = f"""
        Solving equations (i), (ii) and  (iii), 

        $\\displaystyle 
        \\begin{{aligned}}
        
        & {latex(I0)}  = {latex(I0_mod_val)} = {latex(I0_val)}
        \\qquad \\cdots \\text{{ (ii) }}
        
        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())

        soln_stmt = f"""
        
        $\\displaystyle 
        \\begin{{aligned}}
        
        \\implies
        {latex(R)}  
        & = {latex(V / I)}
        \\\\ \\\\
        
        & = \\frac {{ {latex(val_dict[E0])} }}  {{ {latex( abs( ueval(I0_val) )  )} }} 
        \\qquad \\text{{ take absolute value of current as resistance is always positive }}
        \\\\ \\\\
        
        & = {latex(R_val)}
        \\qquad \\cdots \\text{{ (iii) }}

        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())

        
        soln_s = "\n\n".join(soln_stmts)
        
        answer_s = f"""
        $\\displaystyle 
        \\begin{{aligned}}
        
        & {latex(I0)}  = {latex(I0_val)} 
        \\\\ \\\\
        
        & {latex(R)}  = {latex(R_val)}

        \\end{{aligned}}
        $
        """
        
        result = {
            'prob_s': prob_s.strip(),
            'soln_s': soln_s.strip(),
            'answer_s': answer_s.strip(),
            }
        return result

    def _draw_bridge(self, val_dict,
                     E0,
                     R1, R2, R3, R4, R5, 
                     I0, I1, I3, 
                     **kwargs):
        
        import schemdraw
        import schemdraw.elements as elm
        import matplotlib
        matplotlib.use('Agg') 

        dot_color = kwargs.get('dot_color', 'blue')
        show_current = kwargs.get('show_current', True)
        
        R0, R6 = symbols("R0, R6")
        
        labels = {
            E0: f'{latex(val_dict[E0])}',
            R0: f'{latex(I0)}',
            R1: f'{latex(val_dict[R1])}',
            R2: f'{latex(val_dict[R2])}',
            R3: f'{latex(val_dict[R3])}',
            R4: f'{latex(val_dict[R4])}',
            R5: f'{latex(val_dict[R5])}',
            R6: f'{latex(I0)}',
            }

        
        currents = {
            E0: f'',
            R0: f'',
            R1: f'{latex(I1)}',
            R2: f'({latex(I0)} - {latex(I1)})',
            R3: f'{latex(I3)}',
            R4: f'({latex(I0)} - {latex(I3)})',
            R5: f'({latex(I1)} - {latex(I3)})',
            R6: f'',
            }
        
        if show_current:
            label_ltxs = {k:f"{currents[k]} \\; {v}" for k, v in labels.items()}
        else:
            label_ltxs = {k:f"{v}" for k, v in labels.items()}
        
        with schemdraw.Drawing(show=False) as d:
            d.config(inches_per_unit=.3, unit=3, fontsize= 10)
            
            
            elm.Dot().label('A', color = dot_color)
            r0 = elm.Line().label(
                f"${label_ltxs[R0]}$", rotate = True)
            
            elm.Dot().label('B', color = dot_color)
            r1 = elm.Resistor(at = r0.end).theta(30).length(d.unit * 2).label(
                f"${label_ltxs[R1]}$", rotate = True)
            
            elm.Dot().label('C', color = dot_color)
            
            r2 = elm.Resistor(at = r0.end).theta(-30).length(d.unit * 2).label(
                f"${label_ltxs[R2]}$", rotate = True)
            
            elm.Dot().label('D', color = dot_color, loc = 'bottom')
            
            r3 = elm.Resistor(at = r1.end).theta(-30).length(d.unit * 2).label(
                f"${label_ltxs[R3]}$", rotate = True)
            
            elm.Dot().label('E', color = dot_color)
            
            r4 = elm.Resistor(at = r2.end).theta(30).length(d.unit * 2).label(
                f"${label_ltxs[R4]}$", rotate = True)
            
            r5 = elm.Resistor(at = r1.end, to = r2.end).label(
                f"${label_ltxs[R5]}$", rotate = True)
            
            r6 = elm.Line(at = r3.end).right().label(
                f"${label_ltxs[R6]}$", rotate = True)
            
            elm.Dot().label('F', color = dot_color, loc = 'right')
            
            r7 = elm.Line(at = r6.end).down().length(d.unit * 2)
            elm.Dot().label('G', color = dot_color, loc = 'right')

            r9 = elm.Line(at = r0.start).down().length(d.unit * 2)
            elm.Dot().label('K', color = dot_color, loc = 'left')
            
            e8 = elm.SourceV(at = r7.end, to = r9.end).left().label(
                f"${label_ltxs[E0]}$", rotate = True, loc = 'bottom')
            
            
        svg_str = self._convert_d_to_svg_str(d)
        
        return svg_str
        pass

    pass
#class RCircuitHelper:

    