'''
Created on 10-Jun-2021

@author: bthakur
'''
import numpy as np
from xv.util import toIntOrFloat, locatePath
from sympy import Symbol, symbols, Rational, sympify, pi, E, latex, solve, Eq, simplify
import random
import matplotlib.pyplot as plt
import io
from xv.km.util import latex_fraction
import pathlib
import csv
import operator
import inspect
import traceback

from xv.km import SympyUnitSystem


class NumberUnitHelper:
    _NUMBER_SYSTEM_BASE_10_FILE_PATH = "xv_lib_resources/math/number_system.csv"
    _CONVERSION_UNIT_FILE_PATH = "xv_lib_resources/math/conversion_units.csv"
    _IUPAC_PREFIXES_FILE_PATH = "xv_lib_resources/math/iupac_prefixes.csv"


    def __init__(self, **kwargs):
        self.verbose = kwargs.get('verbose', False)

        self.number_system_base_10 = self._init_number_system_base_10()
        self.iupac_prefixes = self._init_iupac_number_prefixes()
        self.conversion_units = self._init_conversion_units()

        self.su_system = SympyUnitSystem(verbose = self.verbose)

        pass


    def calculate_random_dim(self):
        su_system = self.su_system

        dims = su_system.get_common_dimensions_of_physical_quantities()

        quantities = list(dims.keys())
        selected_quantity = random.choice(quantities)
        selected_dim = dims.get(selected_quantity).get('dim', None)

        soln_stmts = []

        prob_s = f"""
        What is dimension of ${selected_quantity}$?
        """

        soln_stmt = f"""
        """
        soln_stmts.append(soln_stmt.strip())

        if selected_dim is None:
            calc_result = su_system.calc_common_dim(selected_quantity, dims)
            if not calc_result is None:
                soln_stmts.extend(calc_result.get('soln_stmts', []))
                selected_dim = calc_result.get('dim', None)
            pass

        if selected_dim is None:
            selected_dim_latex = f"{{ \\text Not Available}}"
            pass
        else:
            selected_dim_latex = f"{latex(selected_dim)}"


        soln_stmt = f"""
        Dimension of $\\displaystyle {latex(selected_quantity)}
        = \left[ {selected_dim_latex} \\right]$
        """
        soln_stmts.append(soln_stmt.strip())


        soln_s = "\n\n".join(soln_stmts)

        answer_s = f"""
        $\\displaystyle \left[ {selected_dim_latex} \\right]$
        """

        result = {
            'quantities': quantities,
            'selected_quantity': selected_quantity,
            'selected_dim': selected_dim,
            'dims': dims,
            'prob_s': prob_s.strip(),
            'soln_s': soln_s.strip(),
            'answer_s': answer_s.strip()
            }

        return result

    def get_conversion_factors_details(self, conversion_factors):
        s1 = ""
        s2 = ""
        factor = 1
        for (u1, n1), (u2, n2) in conversion_factors:
            u1 = u1.replace("-", " ")
            u2 = u2.replace("-", " ")

            u1 = f"\\text {{ {u1} }}"
            u2 = f"\\text {{ {u2} }}"

            n1_frac = latex_fraction(f"{n2}\\;{u2}", f"{n1}\\;{u1}")
            s1 = s1 + f" * $\\displaystyle {n1_frac}$"

            n2_frac = latex_fraction(f"{n2}", f"{n1}")
            s2 = s2 + f" * $\\displaystyle {n2_frac}$"

            factor = factor * n2 / n1
            pass
        return factor, s1, s2
        pass


    def get_conversion_factors_latex(self, conversion_factors):
        cfs1 = []
        cfs2 = []
        factor = 1
        for (u1, n1), (u2, n2) in conversion_factors:
            u1 = u1.replace("-", " ")
            u2 = u2.replace("-", " ")

            n2_u2_by_n1_u1_ltx = latex_fraction(f"{latex(n2)}\\text {{ {u2} }}",
                                                f"{latex(n1)}\\text {{ {u1} }}")

            n2_u2_by_n1_u1 = Symbol(n2_u2_by_n1_u1_ltx)

            cfs1.append(n2_u2_by_n1_u1)


            n2_by_n1_ltx = latex_fraction(f"{latex(n2)}", f"{latex(n1)}")

            n2_by_n1 = Symbol(n2_by_n1_ltx)

            cfs2.append(n2_by_n1)

            factor = factor * n2 / n1
            pass

        cfs1 = [f"{latex(item)}" for item in cfs1]
        cfs1 = f"\\; \\cdot \\;".join(cfs1)

        cfs2 = [f"{latex(item)}" for item in cfs2]
        cfs2 = f"\\; \\cdot \\;".join(cfs2)

        return factor, cfs1, cfs2
        pass


    def get_conversion_factors(self, conversion_set, conversion_path):
        conversion_factors = []
        for idx in range(len(conversion_path) - 1):
            unit1 = conversion_path[idx]
            unit2 = conversion_path[idx + 1]
            conversion_factor  = self.get_conversion_factor(
                                    conversion_set, unit1, unit2)

            if not conversion_factor is None:
                conversion_factors.append(conversion_factor)

            pass

        return conversion_factors
        pass

    def get_conversion_factor(self, conversion_set, unit1, unit2):
        for (u1, n1), (u2, n2) in conversion_set:
            if u1 == unit1 and u2 == unit2:
                return (u1, n1), (u2, n2)

            if u2 == unit1 and u1 == unit2:
                return (u2, n2), (u1, n1)
            pass
        return None
        pass

    def get_conversion_graph(self, conversion_set):
        import networkx as nx
        G = nx.Graph()
        nodes = []
        edges = []

        for (u1, n1), (u2, n2) in conversion_set:
            nodes.append(u1)
            nodes.append(u2)
            edges.append([u1, u2])
            pass

        G.add_nodes_from(nodes)
        G.add_edges_from(edges)

        return G
        pass


    def get_conversion_path(self, conversion_set, u1, u2):
        import networkx as nx
        G = self.get_conversion_graph(conversion_set)
        conv_path = None
        try:
            conv_path = nx.shortest_path(G, source = u1, target = u2)
        except:
            pass

        if not conv_path is None:
            return G, conv_path

        try:
            conv_path = nx.shortest_path(G, source = u2, target = u1)
        except:
            pass

        return G, conv_path
        pass

    def plot_conversion_path(self, G):
        import networkx as nx

        labels = {node:node for node in G.nodes}
        edge_labels = {}
        pos = nx.spring_layout(G)
        fig = plt.figure(num=None, figsize=(10, 6))
        nx.draw(G,
                pos,
                edge_color='black',
                width=1,
                linewidths=1,
                node_size=500,
                node_color='pink',
                alpha=0.9,
                labels = labels,
                font_size = 10
                )

        nx.draw_networkx_edge_labels(G,
                pos,
                edge_labels = edge_labels,
                font_color = 'red'
                )
        plt.axis('off')

        io_buf = io.BytesIO()
        plt.savefig(io_buf, format = "svg", bbox_inches='tight')
        plt.close(fig)

        result = io_buf.getvalue()
        result = result.decode('ascii', 'ignore')
        return result
        pass

    def get_number_system_prefix_value(self, prefix, base = 10):

        number_system_prefixes = self.get_number_system_prefixes(
                        base = 10, set_type = None)

        return number_system_prefixes.get(prefix, None)
        pass

    def get_number_system_prefixes(self, base = 10, set_type = None):
        number_system = self.get_number_system(base = base, set_type = None)

        number_system_prefixes = {}
        for k, v in number_system.items():
            units = self.get_unit_names(v)
            for unit in units:
                number_system_prefixes[unit] = k
                pass
            pass
        return number_system_prefixes
        pass

    def get_number_system_str_expr(self, number_system, base = 10):
        s = ""
        for k, v in number_system.items():
            s = s + "\n" + f"$\\displaystyle {{{base}}}^{{{k}}}$" + f"\t = \t {v}"
            pass
        return s
        pass

    def get_unit_names(self, name_str):
        split_text = ' or '
        names = name_str.split(split_text)
        names = [name.strip() for name in names]
        return names
        pass

    def get_number_system(self, base = 10,
                          set_type = None
                          ):
        '''
        set_type = 'small'   #for less than zero exp
                   'large'   #for more than one exp
                   'common'  #exp between -9 to 9
                   None #for all
        '''
        if base == 10:
            return self._get_number_system_base_10(set_type = set_type)
        return None
        pass


    def _get_number_system_base_10(self, set_type = None):
        if set_type is None:
            return self.number_system_base_10

        number_system = {
            k:v
            for (k, v) in self.number_system_base_10.items()
            if (isinstance(k, int))
            }

        if set_type == 'small':
            minKey = -np.inf
            maxKey = -1
            pass
        elif set_type == 'large':
            minKey = 2
            maxKey = +np.inf
            pass
        elif set_type == 'large':
            minKey = -np.inf
            maxKey = +np.inf
            pass
        else:
            minKey = -9
            maxKey = +9
            pass

        number_system_filtered = {
            k:v
            for (k, v) in number_system.items()
            if (k >= minKey and k <= maxKey)
            }
        return number_system_filtered
        pass


    def get_iupac_number_prefixes(self):
        return self.iupac_prefixes

    def _init_number_system_base_10(self):

        file_path = self.__class__._NUMBER_SYSTEM_BASE_10_FILE_PATH
        from xv.util.data import Repository
        found_file_path = Repository.FIND_EXACT_FILE(file_path, verbose = self.verbose)

        if not found_file_path:
            return {}

        number_system = {}
        with open(found_file_path, newline='') as f:
            reader = csv.reader(f)
            for row in reader:
                if len(row) != 2:
                    print(f"Error: number of columns ={len(row)}: {row}")
                    continue
                n = toIntOrFloat(float(row[0]))
                prefix = row[1]
                number_system[n] = prefix
                pass
            pass
        return number_system
        pass

    def get_random_conversion_category(self):
        return random.choice(list(self.get_conversion_units().keys()))

    def get_units(self, category):
        conversion_set = self.get_conversion_units()[category]

        units = []
        for (u1, n1), (u2, n2) in conversion_set:
            units.append(u1)
            units.append(u2)
            pass
        return list(set(units))
        pass

    def get_conversion_set(self,
                           category,
                           units = None,
                           output_type = 'list',
                           sep = "\t=\t",
                           line = "\n"
                           ):
        '''
        output_type = 'list', or 'list_str', or 'str'
        '''
        conversion_set = self.get_conversion_units()[category]

        if not units is None:
            conversion_set = [
                item
                for item in conversion_set
                if (item[0][0] in units or
                    item[1][0] in units)
                ]
            pass

        if output_type == 'list':
            return conversion_set

        #max_len = max([len(f"{n1} {u1}")
        #               for (u1, n1), (u2, n2) in conversion_set])

        conversion_set_list_str = []
        for (u1, n1), (u2, n2) in conversion_set:
            conversion_set_list_str.append(
                f"{n1} {u1} {sep} {n2} {u2}"
                )
            pass
        if output_type == 'list_str':
            return conversion_set_list_str
            pass

        return line.join(conversion_set_list_str)
        pass

    def get_conversion_units(self):
        return self.conversion_units
        pass


    def _init_iupac_number_prefixes(self):

        file_path = self.__class__._IUPAC_PREFIXES_FILE_PATH
        from xv.util.data import Repository
        found_file_path = Repository.FIND_EXACT_FILE(file_path, verbose = self.verbose)

        if not found_file_path:
            return {}

        _iupac_prefixes = {}
        with open(found_file_path, newline='') as f:
            reader = csv.reader(f)
            for row in reader:
                if len(row) != 2:
                    print(f"Error: number of columns ={len(row)}: {row}")
                    continue
                n = toIntOrFloat(float(row[0]))
                prefix = row[1]
                _iupac_prefixes[n] = prefix
                pass
            pass
        return _iupac_prefixes
        pass

    def _init_conversion_units(self):

        file_path = self.__class__._CONVERSION_UNIT_FILE_PATH
        from xv.util.data import Repository
        found_file_path = Repository.FIND_EXACT_FILE(file_path, verbose = self.verbose)

        if not found_file_path:
            return {}

        cus = {}
        with open(found_file_path, newline='') as f:
            reader = csv.reader(f)
            for row in reader:
                if len(row) != 5:
                    print(f"Error: number of columns ={len(row)}: {row}")
                    continue
                category, u1, n1, u2, n2 = row
                n1 = toIntOrFloat(float(n1))
                n2 = toIntOrFloat(float(n2))

                if not category in cus.keys():
                    cus[category] = []

                cus[category].append([(u1, n1),
                                      (u2, n2)
                                      ])
                pass
            pass

        cus = self._extend_conversion_units(cus)
        return cus
        pass


    def _extend_conversion_units(self, cus):
        cus['Length'].extend(self._extend_conversion_units_Length())
        cus['Mass'].extend(self._extend_conversion_units_Mass())
        cus['Area'].extend(self._extend_conversion_units_Area(cus['Length']))
        cus['Volume'].extend(self._extend_conversion_units_Volume(cus['Length']))
        return cus
        pass

    def _extend_conversion_units_Length(self):
        cus = []
        base = 10
        unit_pair = self._get_derived_conversion_units('meter', 'kilo-', base = base)
        cus.append(unit_pair)

        unit_pair = self._get_derived_conversion_units('meter', 'centi-', base = base)
        cus.append(unit_pair)

        return cus

    def _extend_conversion_units_Mass(self):
        cus = []
        base = 10
        unit_pair = self._get_derived_conversion_units('gram', 'kilo-', base = base)
        cus.append(unit_pair)

        return cus

    def _extend_conversion_units_Area(self, cusLength):
        cus = []
        for (u1, n1), (u2, n2) in cusLength:
            cus.append([(f"square-{u1}", toIntOrFloat(n1**2, 6)),
                        (f"square-{u2}", toIntOrFloat(n2**2, 6))])
            pass
        return cus

    def _extend_conversion_units_Volume(self, cusLength):
        cus = []
        base = 10
        unit_pair = self._get_derived_conversion_units('liter', 'milli-', base = base)
        cus.append(unit_pair)

        for (u1, n1), (u2, n2) in cusLength:
            cus.append([(f"cubic-{u1}", toIntOrFloat(n1**3, 6)),
                        (f"cubic-{u2}", toIntOrFloat(n2**3, 6))])
            pass
        return cus

    def _get_derived_conversion_units(self, unit, prefix, base = 10):
        unit_exp = self.get_number_system_prefix_value(prefix, base = base)
        unit_pair = [(f'{prefix}{unit}', 1),
                    (f'{unit}', 1 * (base ** unit_exp) ) ]
        return unit_pair
        pass



    def get_problem_to_convert_units(self, *args, **kwargs):
        nuh = self
        right_arrow = symbols(u'\U00002192')
        qty =  random.randint(1, 10)

        conversion_category = nuh.get_random_conversion_category()
        conversion_set = nuh.get_conversion_set(conversion_category)
        conversion_set_str = nuh.get_conversion_set(
                            conversion_category, output_type = 'str')
        units = nuh.get_units(conversion_category)

        qu1_org, qu2_org = random.sample(units, 2)

        conversion_graph, conv_path = nuh.get_conversion_path(conversion_set, qu1_org, qu2_org)
        svg_str = nuh.plot_conversion_path(conversion_graph)

        qu1 = Symbol(f"\\text {{ {qu1_org} }}")
        qu2 = Symbol(f"\\text {{ {qu2_org} }}")


        prob_s = f"""
            Convert ${latex(qty)} {latex(qu1)}$ to ${latex(qu2)}$.

            Note: You may use the following table:
            {conversion_set_str}
            """

        soln_stmts = []
        answer_stmts = []
        soln_stmt = f"""
            ${latex(qty)} {latex(qu1)} = \\text {{ ? }} {latex(qu2)}$
            """
        soln_stmts.append(soln_stmt)


        if conv_path is None:
            soln_exception = f"""
                The conversion data is not sufficient
                """
        else:
            soln_exception = None
            conv_path_str = f" {right_arrow} ".join(conv_path)
            soln_stmt = f"""
                The conversion path will be:
                {conv_path_str}
                """
            soln_stmts.append(soln_stmt)

            conversion_factors = nuh.get_conversion_factors(conversion_set, conv_path)

            factor, cfs1, cfs2  = nuh.get_conversion_factors_latex(conversion_factors)

            n_result = factor * qty


            soln_stmt = f"""$ \\displaystyle

            \\begin{{aligned}}

            & {latex(qty)} {latex(qu1)}
            \\\\ \\\\

            & = {latex(qty)} {latex(qu1)} \\cdot {cfs1}
            \\\\ \\\\

            & = {latex(qty)} \\cdot {cfs2} \\cdot {latex(qu2)}
            \\\\ \\\\

            & = {latex(qty)} \\cdot {latex(factor)} \\cdot {latex(qu2)}
            \\\\ \\\\

            & = {latex(n_result)} {latex(qu2)}

            \\end{{aligned}}
            $
            """

            soln_stmts.append(soln_stmt.strip())


            answer_s = f"""$ \\displaystyle
            {latex(n_result)} {latex(qu2)}
            $"""

            answer_stmts.append(answer_s.strip())
            pass

        answer_stmts = "\n\n".join(answer_stmts)
        soln_s_part1 = soln_stmts[0]
        soln_s_part2 = "\n\n".join(soln_stmts[1:])

        result = {
            'prob_s': prob_s,
            'svg_str': svg_str,
            'conv_path': conv_path,
            'soln_exception': soln_exception,
            'soln_s_part1': soln_s_part1,
            'soln_s_part2': soln_s_part2,
            'answer_s': answer_stmts
            }
        return result



    pass
#class NumberUnitHelper