'''
Created on 12-Jul-2021

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

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

# _vr_circuit_helper
# VrCircuitHelper

from ._base_circuit_helper import BaseCircuitHelper

class VrCircuitHelper(BaseCircuitHelper):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        pass
    
    def get_problem_circuit_VR_simple(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, 3))
        
                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_series_V_series_R_circuit(**kwargs)
    
            if random_choice == 2:
                return self._get_series_V_parallel_R_circuit(**kwargs)
    
            return self._get_series_V_series_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_series_V_series_R_circuit(self, **kwargs):
        V = Symbol('V')
        I = Symbol('I')
        R = Symbol('R')
        V_unit = Symbol(f'V')
        R_unit = Symbol(f'\\Omega')
        
        V_count = random.randint(2, 3)
        R_count = random.randint(2, 4)
        
        dot_labels = list('ABCDEFGHKLMNPS')[:V_count + R_count + 1 + 2]
        dot_idx = 0

        V_labels = []
        V_vals = []
        V_val_labels = []
        for idx in range(V_count):
            V_val = Symbol(f"{latex(V)}_{{ {idx} }}")
            V_val_label = Symbol(f"{latex(V)}_{{ {dot_labels[dot_idx]} {dot_labels[dot_idx + 1]} }}")
            
            V_labels.append(dot_labels[dot_idx])
            V_vals.append(V_val)
            V_val_labels.append(V_val_label)
            dot_idx = dot_idx + 1
            pass
        
        # ending label of voltages
        V_labels.append(dot_labels[dot_idx])
        
        V_result_label = Symbol(f"{latex(V)}_{{ {V_labels[0]} {V_labels[-1]} }}")
            
        
        R_labels = []
        R_vals = []
        R_val_labels = []
        V_of_R_labels = []
        for idx in range(R_count):
            R_val = Symbol(f"{latex(R)}_{{ {idx} }}")
            R_val_label = Symbol(f"{latex(R)}_{{ {dot_labels[dot_idx]} {dot_labels[dot_idx + 1]} }}")
            V_of_R_label = Symbol(f"{latex(V)}_{{ {dot_labels[dot_idx]} {dot_labels[dot_idx + 1]} }}")
            
            R_labels.append(dot_labels[dot_idx])
            R_vals.append(R_val)
            R_val_labels.append(R_val_label)
            V_of_R_labels.append(V_of_R_label)
            
            dot_idx = dot_idx + 1
            pass


        
        # ending label of resistances
        R_labels.append(dot_labels[dot_idx])
        
        R_result_label = Symbol(f"{latex(R)}_{{ {R_labels[0]} {R_labels[-1]} }}")
        
        oth_labels = [*dot_labels[dot_idx: dot_idx + 3], V_labels[0]]
        
        
        svg_str = self._draw_series_V_series_R_circuit(V_vals, R_vals, I, 
                                                       is_simplified = False)

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

        soln_stmts = []

        svg_str_with_dots = self._draw_series_V_series_R_circuit(V_vals, R_vals, I, 
                                                       dot_labels = dot_labels)
        
        soln_stmt = f"""
        Let us redraw the circuit labeling the end-points of each component.
        {svg_str_with_dots}
        """
        soln_stmts.append(soln_stmt.strip())
        

        dot_labels_V_result = [V_labels[0], *R_labels, *oth_labels[1:-1] ]
        svg_str_V_result_with_dots = self._draw_series_V_series_R_circuit([V_result_label], 
                                                        R_vals, I, 
                                                        dot_labels = dot_labels_V_result)
        
        
        
        soln_stmt = f"""
        Voltage in series produces cumulative effect by simple addition.
        
        $\\displaystyle 
        \\begin{{aligned}}
        
        {latex(V_result_label)} 
        & = {latex(sum(V_val_labels))}
        \\\\ \\\\
        
        & = {latex(sum(V_vals))}
        \\quad \\cdots \\text{{ (i) }}
      
        \\end{{aligned}}
        $
        {svg_str_V_result_with_dots}
        """
        soln_stmts.append(soln_stmt.strip())
        
        
        V_oth_labels = []
        for idx, oth_label in enumerate(oth_labels[:-1]):
            V_oth_label = Symbol(f"{latex(V)}_{{ {oth_label} {oth_labels[idx + 1]} }}")
            V_oth_labels.append(V_oth_label)
            pass
        
        V_zero_label = Symbol(f"{latex(V)}_{{ {oth_labels[0]} {oth_labels[-1]} }}")
        V_oth_vals = 0
        
        soln_stmt = f"""
        $\\displaystyle 
        \\begin{{aligned}}
        
        & {latex(V_zero_label)}  = {latex(sum(V_oth_labels))}
        \\\\ \\\\
        
        \\implies
        & {latex(V_zero_label)}  = {latex(V_oth_vals)}
        \\qquad \\text{{ as there is not source of voltage or resistance in this part }}
        
        \\end{{aligned}}
        $
        
        """
        soln_stmts.append(soln_stmt.strip())
        
        
        V_oth_labels = []
        for idx, oth_label in enumerate(oth_labels[:-1]):
            V_oth_label = Symbol(f"{latex(V)}_{{ {oth_label} {oth_labels[idx + 1]} }}")
            V_oth_labels.append(V_oth_label)
            pass
        V_all_result_label = Symbol(f"{latex(V)}_{{ {oth_labels[-1]} {oth_labels[0]} }}")
        
        V_all_labels = [V_result_label, *[-item for item in V_of_R_labels]]
        
        IR_val_labels = [I * item for item in R_val_labels]
        
        soln_stmt = f"""
        $\\displaystyle 
        \\begin{{aligned}}
        
        \\implies
        & - {latex(V_all_result_label)} 
         = {latex(V_oth_vals)}
        \\qquad \\text{{ as  }} {latex(V_zero_label)} = - {latex(V_all_result_label)}
        \\\\ \\\\
        
        \\implies
        & {latex(V_all_result_label)} 
         = {latex(V_oth_vals)}
        \\\\ \\\\
        
        \\implies
        & {latex(sum(V_all_labels))} 
         = {latex(V_oth_vals)}
        \\qquad \\text{{ as resistances consume voltage, they are shown as negative voltages. }} 
        \\\\ \\\\


        \\implies
        & {latex(V_result_label)} 
         = {latex(sum(V_of_R_labels))}
        \\\\ \\\\
        
        \\implies
        & {latex(V_result_label)} 
         = {latex(sum(IR_val_labels))}
        \\qquad \\text{{ using }} {latex(V)} = {latex(I*R)} 
        \\\\ \\\\
        
        \\implies
        & {latex(V_result_label)} 
         = {latex(I)} \\left( {latex(sum(R_val_labels))} \\right)
        \\\\ \\\\
        
        \\implies
        & {latex(V_result_label)} 
         = {latex(I)} \\left( {latex(sum(R_val_labels))} \\right)
        
        \\end{{aligned}}
        $
        
        
        """
        soln_stmts.append(soln_stmt.strip())
        
        svg_str_simplified = self._draw_series_V_series_R_circuit(V_vals, R_vals, I, 
                                                                  is_simplified = True)
        
        
        
        soln_stmt = f"""
        Therefore, if we had to replace all resistances with one equivalent resistance $
        {latex(R)}$,

        $\\displaystyle 
        \\begin{{aligned}}
        
        {latex(R)} 
        & = {latex(sum(R_val_labels))}
        \\\\ \\\\
        
        & = {latex(sum(R_vals))}
        \\quad \\cdots \\text{{ (ii) }}

        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())
        
        
        soln_stmt = f"""
        {svg_str_simplified}
        
        Therefore, the current ${latex(I)}$ flowing through the circuit ,

        $\\displaystyle 
        \\begin{{aligned}}
        
        {latex(I)} 
        & =  {latex( V / R  )}
        \\\\ \\\\
        
        & =  {latex( sum(V_vals) / sum(R_vals)  )}
        \\quad \\cdots \\text{{ (iii) }}

        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())
        
        
        
        soln_s = "\n\n".join(soln_stmts)
        
        answer_s = f"""
        $\\displaystyle {latex(V)} = {latex(sum(V_vals))} $
        
        $\\displaystyle {latex(R)} = {latex(sum(R_vals))} $
        
        $\\displaystyle {latex(I)} = {latex( sum(V_vals) / sum(R_vals)  )} $
        """
        
        result = {
            'prob_s': prob_s.strip(),
            'soln_s': soln_s.strip(),
            'answer_s': answer_s.strip(),
            }
        return result

    def _get_series_V_parallel_R_circuit(self, **kwargs):
        V = Symbol('V')
        I = Symbol('I')
        R = Symbol('R')
        V_unit = Symbol(f'V')
        R_unit = Symbol(f'\\Omega')
        
        V_count = random.randint(2, 3)
        R_count = random.randint(2, 4)
        
        dot_labels = list('ABCDEFGHKLMNPS')[:V_count + R_count + 1 + 2]
        dot_idx = 0

        V_labels = []
        V_vals = []
        V_val_labels = []
        for idx in range(V_count):
            V_val = Symbol(f"{latex(V)}_{{ {idx} }}")
            V_val_label = Symbol(f"{latex(V)}_{{ {dot_labels[dot_idx]} {dot_labels[dot_idx + 1]} }}")
            
            V_labels.append(dot_labels[dot_idx])
            V_vals.append(V_val)
            V_val_labels.append(V_val_label)
            dot_idx = dot_idx + 1
            pass
        
        # ending label of voltages
        V_labels.append(dot_labels[dot_idx])
        
        V_result_label = Symbol(f"{latex(V)}_{{ {V_labels[0]} {V_labels[-1]} }}")
            
        
        R_labels = []
        R_vals = []
        I_vals = []
        R_val_labels = []
        V_of_R_labels = []
        V_label_by_R_labels = []
        V_by_R_labels = []
        for idx in range(R_count):
            R_val = Symbol(f"{latex(R)}_{{ {idx} }}")
            I_val = Symbol(f"{latex(I)}_{{ {idx} }}")
            R_val_label = Symbol(f"{latex(R)}_{{ {dot_labels[dot_idx]} {dot_labels[dot_idx + 1]} }}")
            
            next_label = dot_labels[dot_idx + 1]
            V_of_R_label = Symbol(f"{latex(V)}_{{ {next_label}_0 {next_label}_1 }}")
            
            V_label_by_R_label = Symbol(f"{latex(V_of_R_label / R_val)}")
            V_by_R_label = Symbol(f"{latex(V / R_val)}")
            
            
            R_labels.append(dot_labels[dot_idx])
            V_of_R_labels.append(V_of_R_label)
            R_vals.append(R_val)
            I_vals.append(I_val)
            R_val_labels.append(R_val_label)
            V_label_by_R_labels.append(V_label_by_R_label)
            V_by_R_labels.append(V_by_R_label)
            
            dot_idx = dot_idx + 1
            pass

        
        # ending label of resistances
        R_labels.append(dot_labels[dot_idx])
        
        R_result_label = Symbol(f"{latex(R)}_{{ {R_labels[0]} {R_labels[-1]} }}")
        
        oth_labels = [*dot_labels[dot_idx: dot_idx + 3], V_labels[0]]
        
        
        svg_str = self._draw_series_V_parallel_R_circuit(V_vals, R_vals, I, 
                                                       is_simplified = False)

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

        soln_stmts = []

        svg_str_with_dots = self._draw_series_V_parallel_R_circuit(V_vals, R_vals, I, 
                                                       dot_labels = dot_labels)
        
        soln_stmt = f"""
        Let us redraw the circuit labeling the end-points of each component.
        {svg_str_with_dots}
        """
        soln_stmts.append(soln_stmt.strip())
        

        dot_labels_V_result = [V_labels[0], *R_labels, *oth_labels[1:-1] ]
        svg_str_V_result_with_dots = self._draw_series_V_parallel_R_circuit([V_result_label], 
                                                        R_vals, I, 
                                                        dot_labels = dot_labels_V_result)
        
        
        
        soln_stmt = f"""
        Voltage in series produces cumulative effect by simple addition.
        
        $\\displaystyle 
        \\begin{{aligned}}
        
        {latex(V_result_label)} 
        & = {latex(sum(V_val_labels))}
        \\\\ \\\\
        
        & = {latex(sum(V_vals))}
        \\quad \\cdots \\text{{ (i) }}
      
        \\end{{aligned}}
        $
        {svg_str_V_result_with_dots}
        """
        soln_stmts.append(soln_stmt.strip())
        
        V_of_R_labels_latex = " = ".join([f"{latex(item)}" for item in V_of_R_labels])
        
        inv_R_vals = [1/item for item in R_vals]

        R_result_inv_symb = Symbol(f"{latex(sum(inv_R_vals))}")
        
        R_result_symb = Symbol(f"\\frac{{ 1 }}{{ {latex(sum(inv_R_vals))} }}")
        
        
        soln_stmt = f"""
        $\\displaystyle 
        \\begin{{aligned}}
        
        & {latex(I)}  = {latex(sum(I_vals))}
        \\qquad \\text{{ applying Kirchhoff's Laws at point }} {latex(R_labels[1])}_0
        \\\\ \\\\
        
        \\implies
        & {latex(V / R)}  = {latex(sum(V_label_by_R_labels))}
        \\qquad \\text{{ if we replace all resistances with one resistance with }} {latex(R)} 
        \\\\ \\\\
        
        \\implies
        & {latex(V / R)}  = {latex(sum(V_by_R_labels))}
        \\qquad \\text{{ as }} {V_of_R_labels_latex} = {latex(V)}
        \\\\ \\\\
        
        \\implies
        & {latex(1 / R)}  = {latex(R_result_inv_symb)}
        \\\\ \\\\
        
        \\implies
        & {latex(R)}  = {latex(R_result_symb)}
        \\quad \\cdots \\text{{ (ii) }}
        
        \\end{{aligned}}
        $
        
        """
        soln_stmts.append(soln_stmt.strip())
        
        
        svg_str_simplified = self._draw_series_V_series_R_circuit(V_vals, [R_result_symb], I, 
                                                                  is_simplified = True)

        I_symb = Symbol(f"""\\left( {latex( sum(V_vals) )} \\right) 
                            \\left( {latex(R_result_inv_symb)} \\right)""")
        
        soln_stmt = f"""
        {svg_str_simplified}
        
        Therefore, the current ${latex(I)}$ flowing through the circuit ,

        $\\displaystyle 
        \\begin{{aligned}}
        
        {latex(I)} 
        & =  {latex( V / R_result_symb  )}
        \\\\ \\\\
        
        & =  {latex( sum(V_vals) / R_result_symb )}
        \\\\ \\\\

        & = {latex(I_symb)}
        \\quad \\cdots \\text{{ (iii) }}

        \\end{{aligned}}
        $
        """
        soln_stmts.append(soln_stmt.strip())
        
        
        
        soln_s = "\n\n".join(soln_stmts)
        
        answer_s = f"""
        $\\displaystyle {latex(V)} = {latex(sum(V_vals))} $
        
        $\\displaystyle {latex(R)} = \\frac{{ 1 }}{{ {latex(sum(inv_R_vals))} }} $
        
        $\\displaystyle {latex(I)} = {latex(I_symb)} $
        """
        
        result = {
            'prob_s': prob_s.strip(),
            'soln_s': soln_s.strip(),
            'answer_s': answer_s.strip(),
            }
        return result

    def _draw_series_V_series_R_circuit(self, Vs, Rs, I,  **kwargs):
        
        is_simplified = kwargs.get('is_simplified', False)
        dot_labels = kwargs.get('dot_labels', None)
        dot_color = kwargs.get('dot_color', 'blue')
        I_color = kwargs.get('I_color', 'red')
        
        
        label_idx = 0
        
        if is_simplified:
            Vs = [sum(Vs)]
            Rs = [sum(Rs)]
        
        from xv.km.ui import SvgHelper  
        import schemdraw
        import schemdraw.elements as elm
        import matplotlib
        matplotlib.use('Agg') 
        
        dot_radius = 0.001
        

        d = schemdraw.Drawing()
        std_length = d.unit * .75

        if dot_labels:
            d.add(elm.Dot(radius = dot_radius).label(dot_labels[label_idx], 
                                                     color = dot_color))
            label_idx = label_idx + 1
        
        for V_temp in Vs:
            V1 = d.add(elm.SourceV(label=f'${latex(V_temp)}$', d = "right"))
            if dot_labels:
                d.add(elm.Dot(radius = dot_radius).label(dot_labels[label_idx], 
                                                         color = dot_color))
                label_idx = label_idx + 1

        for R_temp in Rs:
            R1 = d.add(elm.Resistor(label=f'${latex(R_temp)}$', d = "right"))
            if dot_labels:
                d.add(elm.Dot(radius = dot_radius).label(dot_labels[label_idx], 
                                                         color = dot_color))
                label_idx = label_idx + 1
        
        d.add(elm.Arrow(arrow='->', d='down', l=std_length).label(f'${latex(I)}$', 
                                                                    color = I_color))

        if dot_labels:
            d.add(elm.Dot(radius = dot_radius).label(dot_labels[label_idx], 
                                                     color = dot_color, 
                                                     loc ='right'))
            label_idx = label_idx + 1

        length = d.unit * (len(Vs) + len(Rs))
        d.add( elm.Arrow(arrow='->', d='left', l=length ).label(f'${latex(I)}$', 
                                                                color = I_color,
                                                                loc ='bottom'  ))
        if dot_labels:
            d.add(elm.Dot(radius = dot_radius).label(dot_labels[label_idx], 
                                                     color = dot_color, 
                                                     loc ='left'))
            label_idx = label_idx + 1

        d.add(elm.Arrow(arrow='->', d='up', l=std_length).label(f'${latex(I)}$', 
                                                                  color = I_color))
        
        image_bytes = d.get_imagedata('svg')
        svg_str = image_bytes.decode('ascii', 'ignore')
        
        svg_helper = SvgHelper(verbose = self.verbose)
        svg_str = svg_helper.remove_unwanted_defs(svg_str)
        return svg_str
        pass

    def _draw_series_V_parallel_R_circuit(self, Vs, Rs, I,  **kwargs):
        
        is_simplified = kwargs.get('is_simplified', False)
        dot_labels = kwargs.get('dot_labels', None)
        dot_color = kwargs.get('dot_color', 'blue')
        I_color = kwargs.get('I_color', 'red')
        
        label_idx = 0
        
        if is_simplified:
            Vs = [sum(Vs)]
            Rs = [1 / sum([1/item for item in Rs])]
        
        from xv.km.ui import SvgHelper  
        import schemdraw
        import schemdraw.elements as elm
        import matplotlib
        matplotlib.use('Agg') 
        
        dot_radius = 0.001
        

        d = schemdraw.Drawing()
        std_length = d.unit * .5

        if dot_labels:
            d.add(elm.Dot(radius = dot_radius
                          ).label(dot_labels[label_idx], 
                                  color = dot_color))
            label_idx = label_idx + 1
        
        for V_idx, V_temp in enumerate(Vs):
            d.add(elm.SourceV(label=f'${latex(V_temp)}$', d = "right"))

            if V_idx == len(Vs) - 1: 
                V_last = d.add(elm.Arrow(arrow='->', d='right', l=std_length, 
                                          ).label(f'${latex(I)}$', 
                                                  color = I_color))

            if dot_labels:
                if V_idx < len(Vs) - 1: 
                    loc ='top'
                else:
                    loc ='right'
                d.add(elm.Dot(radius = dot_radius
                              ).label(dot_labels[label_idx], 
                                      color = dot_color,
                                      loc = loc))
                label_idx = label_idx + 1
                pass
            
            pass
            

        for R_idx, R_temp in enumerate(Rs):
            d.push()
            L_length = std_length * (R_idx + 1)
            
            L_up = d.add(elm.Arrow(arrow='->', d='up', l=L_length, start = V_last.end,
                                   ))

            if dot_labels:
                d.add(elm.Dot(radius = dot_radius
                              ).label(f"${dot_labels[label_idx]}_0$", 
                                      color = dot_color,
                                      loc = 'left'))
        
            L_right = d.add(elm.Arrow(arrow='->', d='right', l=std_length, start = L_up.end,
                                      ).label(f'${latex(I)}_{{{R_idx}}}$', 
                                              color = I_color))
        
            R1 = d.add(elm.Resistor(d = "right", start = L_right.end
                                    ).label(f'${latex(R_temp)}$', 
                                            ))

            L_right = d.add(elm.Arrow(arrow='->', d='right', l=std_length, start = L_up.end,
                                      ).label(f'${latex(I)}_{{{R_idx}}}$', 
                                              color = I_color))
            if dot_labels:
                d.add(elm.Dot(radius = dot_radius
                              ).label(f"${dot_labels[label_idx]}_1$", 
                                      color = dot_color,
                                      loc = 'right'))
                label_idx = label_idx + 1

            d.add(elm.Arrow(arrow='->', d='down', l=L_length, start = L_right.end
                            ))
            
            if R_idx < len(Rs) - 1:
                d.pop()
            pass
            
        d.add(elm.Arrow(arrow='->', d='down', l=std_length
                        ).label(f'${latex(I)}$', 
                                color = I_color))

        if dot_labels:
            d.add(elm.Dot(radius = dot_radius
                          ).label(dot_labels[label_idx], 
                                  color = dot_color, 
                                  loc ='right'))
            label_idx = label_idx + 1

        length = d.unit * len(Vs) + d.unit * 1 + std_length * 3
        d.add( elm.Arrow(arrow='->', d='left', l=length 
                         ).label(f'${latex(I)}$', 
                                 color = I_color,
                                 loc ='bottom'  ))
        if dot_labels:
            d.add(elm.Dot(radius = dot_radius
                          ).label(dot_labels[label_idx], 
                                  color = dot_color, 
                                   loc ='left') )
            
            label_idx = label_idx + 1

        d.add(elm.Arrow(arrow='->', d='up', l=std_length
                        ).label(f'${latex(I)}$', 
                                color = I_color))
        
        image_bytes = d.get_imagedata('svg')
        svg_str = image_bytes.decode('ascii', 'ignore')
        
        svg_helper = SvgHelper(verbose = self.verbose)
        svg_str = svg_helper.remove_unwanted_defs(svg_str)
        return svg_str
        pass


    pass
#class VrCircuitHelper:

    