
'''
Created on 02-Mar-2021

@author: bthakur

'''
from xv.km.kelements import (KPrintElement)
from ...util import MathUtil
from xv.km.kelements import KPrintHelper
from xv.km.util import (
                     latex_fraction, 
                     latex_color,
                     latex_strike,
                     color_cols_of_matrix,
                     strike_cols_of_matrix,
                     multiply_rows_of_matrix,
                     wrap_latex_block, 
                     wrap_latex_statement)

from sympy import latex, simplify
        
from ...util import MathUtil



import os 
from sympy import (latex, Rational, Eq, solve)
from sympy.core.expr import UnevaluatedExpr
import sympy as sp

        
class MatrixSolver():
    def __init__(self, A, X, B, **kwargs):
        self.verbose = kwargs.get('verbose', False)
        self.A = A 
        self.X = X 
        self.B = B 
        self.eqs = None
        self._solution = self._initSolution()
        pass

    def _initSolution(self):
        solution = {}
        
        '''
        THIS WAS JUST FOR A TEST
        self._current_solution = []
        self._initSolution_mako_test()
        solution['mako_test'] = self._current_solution.copy()
        '''
        
        self._current_solution = []
        self._initSolution_row_operations()
        solution['row_operations'] = self._current_solution.copy()
        
        self._current_solution = []
        self._initSolution_crammer_rule()
        solution['crammer_rule'] = self._current_solution.copy()
        
        self._current_solution = []
        self._initSolution_row_operation()
        solution['matrix_row_operation'] = self._current_solution.copy()

        return solution
        pass

    def getSolution(self):
        return self._solution
        pass

    def getSolutionTypes(self):
        if not isinstance(self._solution, dict):
            return None
        return self._solution.keys()
        pass

    def _appendSolution(self, data):
        self._current_solution.append(data)
        pass

    def _getEquationsFromM(self, M, X):
        eqs = []
        A = M[:, :-1]
        B = M[:, -1]
        
        for i in range(X.rows):
            expr = (A[i, :] * X)[0]
            eq = Eq(expr , B[i])
            eqs.append(eq)
            pass
        return eqs
        pass

    def _appendEquationsToSolution(self, eqs):
        for eq in eqs:
            self._appendSolution(KPrintElement(eq, 'sympify/html'))
        pass



    def printSolution(self, **kwargs):
        kph = KPrintHelper()
        
        solution_type = kwargs.get("solution_type", None)
        if (not solution_type is None and 
            isinstance(self._solution, dict)):
            solution_item = self._solution.get(solution_type, None)
            if not solution_item is None:
                return kph.print(solution_item)
            pass
        
        return kph.print(self._solution)
    
        pass
    

    def _initSolution_row_operations(self):
        self._appendSolution(KPrintElement(f"\nMethod 1: Equation Operations:\n{'='*50}", 'text/html'))
        self._appendSolution(KPrintElement("Equations", 'text/html'))

        M = self.A.col_insert(self.A.cols, self.B.T)

        self.eqs = self._getEquationsFromM(M, self.X)
        self._appendEquationsToSolution(self.eqs)

        #forward operation
        for i in range(M.rows - 1, 0, -1):
            self._eq_row_op(M, i)

        self._eq_solve_eqs(M)
        pass


    def _eq_row_op(self, M, org_colidx):
        self._appendSolution(KPrintElement(
            f"\n\nMake coeff of {self.X[org_colidx]} equal:\n{'-'*50}", 'text/html'))
        self._eq_row_op_multiple_rows(M, org_colidx)
        pass

    def _eq_solve_eqs(self, M):
        self.eqs = self._getEquationsFromM(M, self.X)
        sol_X = {}
        for idx, eq in enumerate(self.eqs):
            self._appendSolution(KPrintElement(
                f"\n\nSolve equation_{idx}:\n{'-'*50}", 'text/html'))

            self._appendSolution(KPrintElement(latex(eq), 'text/html'))
            if sol_X:
                self._appendSolution(KPrintElement(
                    f"\nSubstituting values from {sol_X}", 'text/html'))
                eq = eq.subs(sol_X)
                self._appendSolution(KPrintElement(f"$\\displaystyle  {latex(eq)} $", 'text/html'))
                
            soln = solve(eq, self.X[idx], dict = True)
            if soln:
                sol_X.update(soln[0])
            self._appendSolution(KPrintElement(sol_X.copy(), 'text/html'))
            
        pass


    def _initSolution_row_operation(self):
        self._appendSolution(KPrintElement(f"\nMethod 3: Row Operations:\n{'='*50}", 'text/html'))
        self._appendSolution(KPrintElement("Join A and B. Call it matrix M", 'text/html'))
        M = self.A.col_insert(self.A.cols, self.B.T)
        
        M_shape = sp.shape(M)
        M_colored = color_cols_of_matrix(M, M_shape[1] - 1)
        self._appendSolution(f"$\\displaystyle M = {latex(M_colored)}$" )
        
        #forward operation
        for i in range(M.rows):
            self._matrix_row_op(M, i)
        
        #reverse operation
        for i in range(M.rows - 1, 0, -1):
            self._matrix_reverse_row_op(M, i)
            
        
        sol_dict = dict(zip(list(self.X),  list(M[:, -1])))
        ### PrintElement.getLatex(func)
        self._appendSolution(KPrintElement(sol_dict, 'sympy_latex/html'))
        
        pass

    def _eq_row_op_multiple_rows(self, M, org_colidx):

        subM = M
        colidx = org_colidx
                
        col_data = list(subM[:org_colidx + 1, colidx])
        self._appendSolution(KPrintElement(
            f"Coeffs of {self.X[org_colidx]} = {col_data}", 
            'text/html'))
        
        lcm_col = MathUtil.lcm(*col_data)
        self._appendSolution(KPrintElement(
            f"LCM of {col_data} = {lcm_col}. We will make coeffs of {self.X[org_colidx]} equal to {lcm_col}", 
            'text/html'))

        row_operands = [int(lcm_col/itm) for itm in col_data]
        self._appendSolution(KPrintElement(
            f"Multiply equations with {row_operands}", 
            'text/html'))
        
        for rowid, row_operand in enumerate(row_operands):
            if rowid > org_colidx:
                continue
            subM.row_op(rowid, lambda v, j:  v * row_operand )

        self.eqs = self._getEquationsFromM(subM, self.X)
        self._appendEquationsToSolution(self.eqs)

            
        self._appendSolution(KPrintElement(
            f"\nDeduct equation_{org_colidx} from others", 
            'text/html'))
        
        for rowid, row_operand in enumerate(row_operands):
            if rowid >= org_colidx:
                break
            subM.row_op(rowid, lambda v, j:  v - subM[org_colidx, j] )
            
        self.eqs = self._getEquationsFromM(subM, self.X)
        self._appendEquationsToSolution(self.eqs)

        
        self._appendSolution(KPrintElement(
            f"\nSimplify equations by dividing each equation with HCF of its coeffs", 
            'text/html'))
        
        for rowid, row_operand in enumerate(row_operands):
            row_data = list(subM[rowid, :])
            row_hcf = MathUtil.hcf(*row_data)

            self._appendSolution(KPrintElement(
                f"HCF coeffs of equation_{rowid} of {[row_data]} = {row_hcf}", 
                'text/html'))
            
            if rowid >= org_colidx:
                #break
                pass
            subM.row_op(rowid, lambda v, j:  int(v / row_hcf) )
            
        self.eqs = self._getEquationsFromM(subM, self.X)
        self._appendEquationsToSolution(self.eqs)

        pass



    def _matrix_reverse_row_op(self, M, org_colidx):
        self._appendSolution(KPrintElement(f"\n\nProcessing reverse column_{org_colidx}:\n{'-'*50}", 'text/html'))
        
        colidx = org_colidx - 1
        while colidx > -1:
            cell_operand = M[colidx, org_colidx]
    
            self._appendSolution(KPrintElement(
                f"Multiply  row_{org_colidx} by {cell_operand} and deduct from row_{colidx} ", 
                'text/html'))
            
            M.row_op(colidx, lambda v, j:  Rational(v - cell_operand * M[org_colidx, j] ))
    
            M_shape = sp.shape(M)
            M_colored = color_cols_of_matrix(M, M_shape[1] - 1)
            stmt = f""" $\\displaystyle 
            M
            = {latex(M_colored)} 
            $ """
            self._appendSolution(stmt)
            

            colidx = colidx - 1

        pass

    def _matrix_row_op(self, M, org_colidx):
        stmt = KPrintElement(f"\nProcessing column_{org_colidx}:\n{'-'*50}", 'text/html')
        
        self._appendSolution(stmt)

        if org_colidx != M.rows - 1:
            self._matrix_row_op_multiple_rows(M, org_colidx)
                        
        self._matrix_row_op_normalize(M, org_colidx)
        pass


    def _matrix_row_op_normalize(self, M, org_colidx):
        cell_operand = M[org_colidx, org_colidx]
        self._appendSolution(KPrintElement(
            f"Divide row_{org_colidx} by {cell_operand}", 
            'text/html'))
        
        M.row_op(org_colidx, lambda v, j:  Rational(v / cell_operand))

        M_shape = sp.shape(M)
        M_colored = color_cols_of_matrix(M, M_shape[1] - 1)
        stmt = f""" $\\displaystyle 
        M
        = {latex(M_colored)} 
        $ """
        self._appendSolution(stmt)
        
        pass


    def _matrix_row_op_multiple_rows(self, M, org_colidx):
        if org_colidx == M.rows - 1:
            return 
        
        
        subM_org = M[org_colidx:, org_colidx:]
        subM = M[org_colidx:, org_colidx:]
        colidx = 0

        #m_str = f"subM_{org_colidx} = {latex(subM)}" 
        subM_latex = f"$\\displaystyle subM_{org_colidx} = " + KPrintElement.getLatex(subM) + "$"
        self._appendSolution(KPrintElement(subM_latex, 'text/html'))

        col_data = list(subM[:, colidx])
        self._appendSolution(KPrintElement(
            f"column_{colidx} = {col_data}", 
            'text/html'))
        
        lcm_col = MathUtil.lcm(*col_data)
        self._appendSolution(KPrintElement(
            f"""LCM of {col_data} = {lcm_col}. 
                We will make each element of column_{colidx} equal to {lcm_col}
            """.strip(), 
            'text/html'))
        
        row_operands = [int(lcm_col/itm) for itm in col_data]
        self._appendSolution(KPrintElement(
            f"Multiply rows with {row_operands}", 
            'text/html'))
        
        for rowid, row_operand in enumerate(row_operands):
            subM.row_op(rowid, lambda v, j:  v * row_operand )


        subM_shape = sp.shape(subM)
        subM_colored = color_cols_of_matrix(subM, subM_shape[1] - 1)
        stmt = f""" $\\displaystyle 
        subM_{org_colidx} 
        = {latex(multiply_rows_of_matrix(subM_org, row_operands))} 
        = {latex(subM_colored)} 
        $ """
        self._appendSolution(stmt)
            
        self._appendSolution(KPrintElement(
            f"Deduct  row_{colidx} from other rows", 
            'text/html'))
        
        for rowid, row_operand in enumerate(row_operands):
            if rowid == colidx:
                continue
            subM.row_op(rowid, lambda v, j:  v - subM[colidx, j] )
            
        #m_str = f"subM_{org_colidx} = {latex(subM)}" 
        subM_latex = f"$\\displaystyle subM_{org_colidx} = " + KPrintElement.getLatex(subM) + "$"
        self._appendSolution(KPrintElement(subM_latex, 'text/html'))



        self._appendSolution(KPrintElement("Updating matrix M", 'text/html'))
        
        for i in range(subM.rows):
            for j in range(subM.cols):
                M[i + org_colidx, j + org_colidx] = subM[i, j]
            pass

        M_shape = sp.shape(M)
        M_colored = color_cols_of_matrix(M, M_shape[1] - 1)
        stmt = f""" $\\displaystyle 
        M
        = {latex(M_colored)} 
        $ """
        self._appendSolution(stmt)
        pass

    
    def _initSolution_crammer_rule(self):
        stmt = f"<h3>Method : Crammer's Rule:</h3>"
        self._appendSolution(KPrintElement(stmt, 'text/html'))

        M = self.A.col_insert(self.A.cols, self.B.T)
        M_id = "M"
        A = self.A
        A_id = "A"
        D = A.det()
        D_id = "D"

        M_shape = sp.shape(M)
        M_colored = color_cols_of_matrix(M, M_shape[1] - 1)
        
        col_den = list(range(self.A.cols))
        
        num_matrixes = []
        num_determinants = []
        stmt_matrixes = []
        stmt_determinants = []
        var_values_str = []
        for i in col_den:
            col_i = [j for j in col_den if j != i]
            col_i.insert(0, -1)

            var_name = f"{self.X[i]}"
            D_i_id = f"D_{{{var_name}}}"
            M_i_id = f"M_{{{var_name}}}"
            M_i = M[:, col_i]
            D_i = M_i.det()
            
            num_matrixes.append([M_i_id, M_i])
            num_determinants.append([D_i_id, D_i])

            stmt_matrix = f"""
            Leave coefficient of {var_name}, taking columns {col_i} from M:
            
            $\\displaystyle 
            {D_i_id}
            = det\\left({M_i_id}\\right) 
            = det\\left({latex(strike_cols_of_matrix(M_colored, i))}\\right)
            = det\\left({latex(M_i)}\\right)
            = {latex(D_i)}
            $
            """.strip()
            stmt_matrixes.append(stmt_matrix)

            det_frc_id  = latex_fraction(D_i_id, D_id)
            det_frc = latex_fraction(latex(D_i), latex(D))
            sign_val = UnevaluatedExpr(-1)**UnevaluatedExpr(i)
            sign_val_n = simplify(sign_val)
            
            var_val = Rational(sign_val_n * D_i, D)
            
            stmt_determinant = f"""
            $\\displaystyle 
            {var_name}
            = {latex(sign_val)} * {det_frc_id}
            = {sign_val_n} * {det_frc}
            = {latex(var_val)}
            $
             
            """
            
            var_value_str = f"""
            $\\displaystyle {var_name} = {latex(var_val)}$
            """
            var_values_str.append(var_value_str.strip())
            
            stmt_determinants.append(stmt_determinant)
            
            pass
        
        stmt_matrixes = "\n".join(stmt_matrixes) 
        stmt_determinants = "\n".join(stmt_determinants)
        var_values_str = ", ".join(var_values_str)
        
        
        stmt1 = f"""
        <h3>Method : Crammer's Rule:</h3>
        Taking all coefficients of variables and result:
        
        $\\displaystyle 
        {M_id}
        = {latex(M_colored)}
        $

        Taking coefficients of all variables, we get denominator matrix A:
        
        $\\displaystyle 
        {D_id}
        = det\\left({A_id}\\right) 
        = det\\left({latex(strike_cols_of_matrix(M, -1))}\\right)
        = det\\left({latex(A)}\\right)
        = {latex(D)}
        $

        {stmt_matrixes}
        
        Compute values of variables:
        
        {stmt_determinants}
        
        Therefore:
        
        {var_values_str}
        """
        
        stmt = f"### This line is not being printed."
        self._appendSolution(KPrintElement(stmt, 'text/html'))
        
        self._appendSolution(KPrintElement(stmt1, 'text/html'))
        pass
    
    

    def _initSolution_mako_test(self):
        self._appendSolution(KPrintElement(f"\nMethod : Mako Test:\n{'='*50}", 'text/html'))
        
        sol_dict = [] 
        data = {'name': 'test_name'}    
        sText = """
        this is _mako_transform_text name = ${name}
        """
        sText = self._mako_transform_text(sText, data)
        sol_dict.append(sText)
        
        
        
        sText = self._mako_transform_template('./templates/test_mako.html', data)
        sol_dict.append(sText)
        
        self._appendSolution(sol_dict)
        pass
    
    def _mako_transform_text(self, text, data):
        from xv.util import MakoHelper
        mako_helper = MakoHelper(verbose = self.verbose)
        rendered_text = mako_helper.render_text(text, **data)
        return rendered_text
        pass

    def _mako_transform_template(self, 
                                 templatePath, 
                                 data,
                                 lookupFolders = None
                                 ):
        '''
         lookupFolders for other resources like include etc.
        '''

        if not os.path.isfile(templatePath):
            dir_path = os.path.dirname(os.path.realpath(__file__))
            templatePath = os.path.join(dir_path, templatePath)
            pass
        
        if not os.path.isfile(templatePath):
            sText = f"""
            file not found
            templatePath = {templatePath}
            data = {data}
            lookupFolders = {lookupFolders}
            """
            print(sText)
            return None
            pass

        from xv.util import MakoHelper
        mako_helper = MakoHelper(verbose = self.verbose)
        rendered_text = mako_helper.render_file(templatePath, 
                                                lookup_folders = lookupFolders, 
                                                **data)
        return rendered_text
        pass

    pass