'''
Created on 20-Jan-2022

@author: bthakur
'''

import random
import inspect
import traceback


from sympy.core.expr import UnevaluatedExpr
from operator import mul
from functools import reduce
from collections import OrderedDict
            
from sympy import ( latex, solve, simplify, 
                    pi, Rational,
                    exp, sin, cos, tan,
                    asin, acos, atan,
                    sinh, cosh, tanh,
                    log, exp, sqrt,
                    Eq
                    )

import numpy as np

from xv.util import list_to_str
from xv.util import run_with_timeout

from xv.km.util import latex_fraction, latex_exponent, wrap_bracket

from .._algebra import BinomialHelper

# _probability_intro_helper
# ProbabilityIntroHelper

class ProbabilityIntroHelper():
    
    def __init__(self, **kwargs):
        self.verbose = kwargs.get('verbose', False)
        self.binomial_helper = BinomialHelper(verbose = self.verbose)
        pass


    def get_random_letter_set(self, count = None):
        population = list("abcdefghijklmnopqrstuvwxyz")
        if count is None:
            count = random.randint(2, 5)
        population = random.sample(population, count)
        
        letters = []
        for i in range(count):
            letter_count = random.choice([1, random.randint(1, 5)])
            letters.extend([population[i]] * letter_count)
            pass
        letters = random.sample(letters, len(letters))
        return letters

    def _latex_of_series(self, tpl, func):
        s_latex = [func(k, v) 
               for k, v in tpl ]
        s_latex = " \\; ".join(s_latex)
        return s_latex
        pass

    def _combination_of_series(self, tpl, func, evaluate = True):
        if evaluate:
            terms = [func(k, v) 
                   for k, v in tpl ]
        else:
            terms = [UnevaluatedExpr(func(k, v)) 
                   for k, v in tpl ]
            
        comb = reduce(mul, terms, 1)
        return comb
        pass

    def form_words_from_letters(self, letters, r = None):
        if r is None or r == len(letters):
            return self._form_words_from_letters(letters)
        
        # for non repeated letters
        letter_counts = {i:letters.count(i) for i in set(letters) 
                         if letters.count(i) > 1}
        if not letter_counts:
            return self._form_words_from_letters(letters)
            
        return self._form_shorter_words_from_letters(letters, r)

    def _get_letter_sets(self, letters, r):
        letter_counts = {i:letters.count(i) for i in set(letters)}
        n = sum([v for v in letter_counts.values()])
        
        letter_sets = []
        letter_sets.append(letter_counts)
        while n > r:
            n = n - 1
            temp_sets = []
            temp_str_sets = set()
            for letter_set in letter_sets:
                for k, v in letter_set.items():
                    if v == 0:
                        continue
                    temp_letter_set = letter_set.copy()
                    temp_letter_set[k] = temp_letter_set[k] - 1 
                    
                    temp_letter_set = OrderedDict(temp_letter_set)
                    if str(temp_letter_set) in temp_str_sets:
                        continue
                    
                    temp_str_sets.add(str(temp_letter_set))
                    
                    temp_sets.append(temp_letter_set)
                    pass
                pass
                
            letter_sets = temp_sets
            pass
        
        return letter_sets
        pass

    def _form_shorter_words_from_letters(self, letters, r):
        if r > len(letters):
            return None
        if r < 0:
            return None
        if r != int(r):
            return None

        bh = self.binomial_helper
        n = len(letters)
        
        letter_counts = {i:letters.count(i) for i in set(letters)}
        repeated_letters = {k:v for k, v in letter_counts.items()
                            if v > 1}
        
        letter_sets = self._get_letter_sets(letters, r)
        
        soln_s = []
        nPrs = []
        answer = 1
        for letter_set in letter_sets:
            item_result = self._form_words_from_letters_short_statement(letter_set)
            soln_s.append(item_result['soln_s'].strip())
            nPrs.append(item_result['nPr'])
            answer = answer + item_result['nPr']
            pass

        nPrs_str = [f"{item}" for item in nPrs]
        nPrs_str = " + ".join(nPrs_str)
        
        stmt = f"""
        As the number of letters({n}) is more then the desired number of letters({r}) in the word,
        we will consider each possible combination of letters and then add the results.
        """
        soln_s.insert(0, stmt.strip())
        
        stmt = f"""
        The number of ways to form words
        $\\displaystyle = {nPrs_str}$
        $\\displaystyle = {answer}$
        """
        soln_s.append(stmt.strip())
        
        soln_s = "\n\n".join(soln_s)
        
        answer_s = f"{answer}"
        
        result = {
            'soln_s': soln_s,
            'answer_s': answer_s,
            }
        return result

    def _form_words_from_letters_short_statement(self, letter_counts):
        bh = self.binomial_helper
        
        letter_counts = {k:v for k, v in letter_counts.items() if v > 0}
        repeated_letters = {k:v for k, v in letter_counts.items()
                            if v > 1}
        
        n = sum([v for v in letter_counts.values()])
        r = n
        
        letter_counts_str = [f"{v}{k}" for k, v in letter_counts.items()]
        letter_counts_str = ", ".join(letter_counts_str)
        
        nPr_latex = bh.get_permutation_latex(n, r)
        nPr_factorial_latex = bh.get_permutation_factorial_latex(n, r)
        nPr_factorial_evaluated_latex = bh.get_permutation_factorial_latex(n, r,
                                                                           evaluate = True)
        nPr = bh.get_permutation_expr(n, r)
        
        soln_s = []
        
        if repeated_letters:
            repeated_letters_str = [f"$\\displaystyle {k}({v}) = {bh.get_permutation_latex(v, v) }$"
                                    for k, v in repeated_letters.items()]
            repeated_letters_str = "\n\n".join(repeated_letters_str)
            
            repeat_v_tuple = [(v, v) for k, v in repeated_letters.items()]
            
            correction_for_repeat = self._latex_of_series(
                repeat_v_tuple, 
                bh.get_permutation_latex)
            
            
            correction_for_repeat_factorial = self._latex_of_series(
                repeat_v_tuple, 
                bh.get_permutation_factorial_latex)
            
            
            corrected_nPr_latex = latex_fraction(nPr_latex, correction_for_repeat)
            corrected_nPr_factorial_latex = latex_fraction(nPr_factorial_latex, correction_for_repeat_factorial)
            
            
            correction_expr_for_repeat = self._combination_of_series(
                repeat_v_tuple, 
                bh.get_permutation_expr)
            
            
            corrected_nPr = Rational(nPr, correction_expr_for_repeat)
            
            stmt = f"""
            Words from {letter_counts_str} = $\\displaystyle 
            = {corrected_nPr_latex}
            = {corrected_nPr_factorial_latex}
            = {latex(corrected_nPr)}
            $
            """
            soln_s.append(stmt.strip())
        else:
            repeated_letters_str = [f"$\\displaystyle {k}: {v} = {bh.get_permutation_latex(v, v) }$"
                                    for k, v in repeated_letters.items()]
            repeated_letters_str = "\n".join(repeated_letters_str)
            
            correction_for_repeat = [bh.get_permutation_latex(v, v) 
                                     for k, v in repeated_letters.items() ]
            
            correction_for_repeat = " \\; ".join(correction_for_repeat)
            
            corrected_nPr = nPr
            
            stmt = f"""
            Words from {letter_counts_str} = $\\displaystyle 
            = {nPr_factorial_latex}
            = {nPr_factorial_evaluated_latex}
            = {latex(corrected_nPr)}
            $
            """
            soln_s.append(stmt.strip())
            pass
            
        
        soln_s = "\n\n".join(soln_s)
        
        result = {
            'nPr_latex': nPr_latex,
            'nPr_factorial_latex': nPr_factorial_latex,
            'nPr_factorial_evaluated_latex': nPr_factorial_evaluated_latex,
            'nPr': corrected_nPr,
            'answer_s': f""" $\\displaystyle {latex(corrected_nPr)}$ """,
            'soln_s': soln_s
            }
        return result


    def _form_words_from_letters(self, letters):
        r = len(letters)
        bh = self.binomial_helper
        
        letter_counts = {i:letters.count(i) for i in set(letters)}
        repeated_letters = {k:v for k, v in letter_counts.items()
                            if v > 1}
        
        n = sum([v for v in letter_counts.values()])
            
        
        nPr_latex = bh.get_permutation_latex(n, r)
        nPr_factorial_latex = bh.get_permutation_factorial_latex(n, r)
        nPr_factorial_evaluated_latex = bh.get_permutation_factorial_latex(n, r,
                                                                           evaluate = True)
        nPr = bh.get_permutation_expr(n, r)
        
        soln_s = []
        
        stmt = f"""
        ways of arranging {r} from {n} items
        
        $\\displaystyle = {nPr_latex}$
        
        """
        soln_s.append(stmt.strip())
        
        if repeated_letters:
            repeated_letters_str = [f"$\\displaystyle {k}({v}) = {bh.get_permutation_latex(v, v) }$"
                                    for k, v in repeated_letters.items()]
            repeated_letters_str = "\n\n".join(repeated_letters_str)
            
            repeat_v_tuple = [(v, v) for k, v in repeated_letters.items()]
            
            correction_for_repeat = self._latex_of_series(
                repeat_v_tuple, 
                bh.get_permutation_latex)
            
            #correction_for_repeat = [bh.get_permutation_latex(v, v) 
            #                         for k, v in repeated_letters.items() ]
            #correction_for_repeat = " \\; ".join(correction_for_repeat)
            
            correction_for_repeat_factorial = self._latex_of_series(
                repeat_v_tuple, 
                bh.get_permutation_factorial_latex)
            
            #correction_for_repeat_factorial = [bh.get_permutation_factorial_latex(v, v) 
            #                         for k, v in repeated_letters.items() ]
            #correction_for_repeat_factorial = " \\; ".join(correction_for_repeat_factorial)
            
            corrected_nPr_latex = latex_fraction(nPr_latex, correction_for_repeat)
            corrected_nPr_factorial_latex = latex_fraction(nPr_factorial_latex, correction_for_repeat_factorial)
            
            
            correction_expr_for_repeat = self._combination_of_series(
                repeat_v_tuple, 
                bh.get_permutation_expr)
            
            #correction_expr_for_repeat = [bh.get_permutation_expr(v, v) 
            #                              for k, v in repeated_letters.items() ]
            
            #correction_expr_for_repeat = reduce(mul, correction_expr_for_repeat, 1)
            
            corrected_nPr = Rational(nPr, correction_expr_for_repeat)
            
            stmt = f"""
            repeated letters and correction for them
            {repeated_letters_str}
            
            Therefore, correction for all repeated letters
            
            $\\displaystyle = {correction_for_repeat}$
            
            Therefore, ways of arranging {r} from {n} items after correction
            
            $\\displaystyle = {corrected_nPr_latex}$
            
            $\\displaystyle = {corrected_nPr_factorial_latex}$
            
            $\\displaystyle = {latex(corrected_nPr)}$
            
            """
            soln_s.append(stmt.strip())
        else:
            repeated_letters_str = [f"$\\displaystyle {k}: {v} = {bh.get_permutation_latex(v, v) }$"
                                    for k, v in repeated_letters.items()]
            repeated_letters_str = "\n".join(repeated_letters_str)
            
            correction_for_repeat = [bh.get_permutation_latex(v, v) 
                                     for k, v in repeated_letters.items() ]
            
            correction_for_repeat = " \\; ".join(correction_for_repeat)
            
            corrected_nPr = nPr
            
            stmt = f"""
            $\\displaystyle = {nPr_factorial_latex}$
            
            $\\displaystyle = {nPr_factorial_evaluated_latex}$
            
            $\\displaystyle = {latex(corrected_nPr)}$
            
            """
            soln_s.append(stmt.strip())
            pass
            
        
        soln_s = "\n\n".join(soln_s)
        
        result = {
            'nPr_latex': nPr_latex,
            'nPr_factorial_latex': nPr_factorial_latex,
            'nPr_factorial_evaluated_latex': nPr_factorial_evaluated_latex,
            'nPr': corrected_nPr,
            'answer_s': f""" $\\displaystyle {latex(corrected_nPr)}$ """,
            'soln_s': soln_s
            }
        return result


    def get_random_chance_data(self, *args, **kwargs):
        item_names1 = ['cake', 'ice cream', 'pizza', 'burger', 'chocolate']
        item_names2 = ['red ball', 'green ball', 'blue ball', 'yellow ball', 
                       'white ball', 'black ball']
        
        item_names = random.choice([item_names1, item_names2])
        
        count = random.randint(2, len(item_names))
        item_names = random.sample(item_names, count)
        item_counts = list(np.random.randint(1, 10, count))
        return item_names, item_counts

    def get_names_counts_str(self, item_names, item_counts):
        s_names_counts = [f"{name}: {item_counts[i]}" 
                        for i, name in enumerate(item_names)]
        
        s_names_counts = "\n".join(s_names_counts)  
        return s_names_counts
    
    
    def get_prob_statement(self, item_names, item_counts, desired_selections):
        
        if len(desired_selections) == 1:
            return self._get_prob_statement_single_item(item_names, 
                                                        item_counts, 
                                                        desired_selections)
            
        return self._get_prob_statement_multiple_items(item_names, 
                                                       item_counts, 
                                                       desired_selections)
        pass

    def _get_prob_statement_multiple_items(self, item_names, item_counts, desired_selections):
        bh = self.binomial_helper
        
        pop_dict = dict(zip(item_names, item_counts))
        selection_k_tuple = [(k, pop_dict[k])
                             for k, v in desired_selections.items()]
        
        selection_v_tuple = [(pop_dict[k], v)
                             for k, v in desired_selections.items()]
        
        desired_selections_str = {f"{v} {k} " for k, v in desired_selections.items()}
        desired_selections_str = ", ".join(desired_selections_str)
        
        total_count = sum(item_counts)
        count = sum(list(desired_selections.values()))
        
        prob_statement_latex = latex_fraction(
            f"\\text{{ways of selecting ({desired_selections_str}) }}",
            f"\\text{{ways of selecting {count} from {total_count} items }}",
                                              )
        
        mCr_latex = [bh.get_binomial_latex(k, v)
                     for k, v in selection_v_tuple]
        mCr_latex = "\\;".join(mCr_latex)
        
        nCr_latex = bh.get_binomial_latex(total_count, count)
        
        prob_latex = latex_fraction(mCr_latex, nCr_latex)
        

        mCr_factorial_latex = self._latex_of_series(
            selection_v_tuple, 
            bh.get_combination_factorial_latex)

        #mCr_factorial_latex = [bh.get_combination_factorial_latex(k, v)
        #             for k, v in selection_v_dict.items()]
        #mCr_factorial_latex = "\\;".join(mCr_factorial_latex)
        
        nCr_factorial_latex = bh.get_combination_factorial_latex(total_count, count)
        
        prob_factorial_latex = latex_fraction(mCr_factorial_latex, nCr_factorial_latex)
        
        mCr_unevaluated = self._combination_of_series(
            selection_v_tuple, 
            bh.get_combination_expr,
            evaluate = False)

        mCr = self._combination_of_series(
            selection_v_tuple, 
            bh.get_combination_expr)


        
            
        #mCr = [bh.get_combination_expr(k, v)
        #             for k, v in selection_v_dict.items()]
        #mCr = reduce(mul, mCr, 1)
        
        nCr = bh.get_combination_expr(total_count, count)
        
        prob_unevaluated = latex_fraction(latex(mCr_unevaluated), latex(nCr))
        prob = Rational(mCr, nCr)
        
        
        result = {
            'total_count': total_count,
            'desired_selections_str': desired_selections_str,
            'prob_statement_latex': prob_statement_latex,
            'mCr_latex': mCr_latex,
            'nCr_latex': nCr_latex,
            'prob_latex': prob_latex,
            'prob_factorial_latex': prob_factorial_latex,
            'prob_unevaluated': prob_unevaluated,
            'prob': prob,
            }
        return result
        pass


    def _get_prob_statement_single_item(self, item_names, item_counts, desired_selections):
        bh = self.binomial_helper
        
        item_name, count = list(desired_selections.items())[0]

        total_count = sum(item_counts)
        item_count = item_counts[item_names.index(item_name)]
        
        prob_statement_latex = latex_fraction(
            f"\\text{{ways of selecting {count} from {item_count} {item_name} }}",
            f"\\text{{ways of selecting {count} from {total_count} items }}",
                                              )
        mCr_latex = bh.get_binomial_latex(item_count, count)
        nCr_latex = bh.get_binomial_latex(total_count, count)
        prob_latex = latex_fraction(mCr_latex, nCr_latex)
        
        
        
        prob_factorial_latex = latex_fraction(
            bh.get_combination_factorial_latex(item_count, count),
            bh.get_combination_factorial_latex(total_count, count)
            )
        
        prob_unevaluated = latex_fraction(bh.get_combination_expr(item_count, count), 
                        bh.get_combination_expr(total_count, count))

        
        prob = Rational(bh.get_combination_expr(item_count, count), 
                        bh.get_combination_expr(total_count, count))
        
        result = {
            'total_count': total_count,
            'item_count': item_count,
            'prob_statement_latex': prob_statement_latex,
            'mCr_latex': mCr_latex,
            'nCr_latex': nCr_latex,
            'prob_latex': prob_latex,
            'prob_factorial_latex': prob_factorial_latex,
            'prob_unevaluated': prob_unevaluated,
            'prob': prob,
            }
        return result
        pass



    def get_random_chance_count(self, *args, **kwargs):
        item_names, item_counts = self.get_random_chance_data()
        
        answer_count = random.choice(item_counts)
        total_count = sum(item_counts)
        answer_name = item_names[item_counts.index(answer_count)]
        
        prob_latex1 = latex_fraction(f"\\text{{Number of {answer_name}}}", 
                                            f"\\text{{Total number of items}}")
        
        prob_latex2 = latex_fraction(answer_count, total_count)
        
        prob = Rational(answer_count, total_count)
        
        s_names_counts = self.get_names_counts_str(item_names, item_counts)

        result = {
            'item_names': item_names,
            'item_counts': item_counts,
            's_names_counts': s_names_counts,
            'answer_count': answer_count,
            'total_count':total_count,
            'answer_name': answer_name,
            'prob_latex1': prob_latex1,
            'prob_latex2': prob_latex2,
            'prob': prob,
            }
        
        return result

    
    def get_single_dice_probability(self, *args, **kwargs):
        population, outcomes = self._get_single_dice_data()
        desired_outcome = random.choice(outcomes)


        prob_latex1 = latex_fraction(f"\\text{{Count of desired outcomes }}", 
                                            f"\\text{{Count of all possible outcomes}}")
        
        prob_latex2 = latex_fraction(len(desired_outcome[1]), 
                                            len(population))
        
        prob = Rational(len(desired_outcome[1]), len(population))
        
        soln_s = f"""
        Possible outcomes: {set(population)}
        Count of possible outcomes: {len(population)}
        
        Desired outcome = {desired_outcome[0]} = {{ {list_to_str(desired_outcome[1])} }}
        Count of desired outcomes: {len(desired_outcome[1])}
        
        Probability({desired_outcome[0]})
        
        $\\displaystyle = {prob_latex1}$
        
        $\\displaystyle = {prob_latex2}$
        
        $\\displaystyle = {latex(prob)}$
        """

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

        result = {
            'population': population,
            'outcomes': outcomes,
            'desired_outcome': desired_outcome,
            'prob_latex1': prob_latex1,
            'prob_latex2': prob_latex2,
            'prob': prob,
            'soln_s': soln_s,
            'answer_s': answer_s,
            }
        
        return result

    
    def _get_single_dice_data(self):
        population = list(range(1, 7))

        outcomes = []
        outcome = ['even numbers', 
                   [x for x in population 
                    if x % 2 == 0]]
        outcomes.append(outcome)
        
        outcome = ['odd numbers', 
                   [x for x in population 
                    if x % 2 != 0]]
        outcomes.append(outcome)
        
        min_n = random.choice(population)
        outcome = [f'more than {min_n}', 
                   [x for x in population 
                    if x > min_n]]
        outcomes.append(outcome)
        
        outcome = [f'more than or equal to {min_n}', 
                   [x for x in population 
                    if x >= min_n]]
        outcomes.append(outcome)
        
        max_n = random.choice(population)
        outcome = [f'less than {max_n}', 
                   [x for x in population 
                    if x < max_n]]
        outcomes.append(outcome)
        
        max_n = random.choice(population)
        outcome = [f'less than or equal to {max_n}', 
                   [x for x in population 
                    if x <= max_n]]
        outcomes.append(outcome)
        
        min_n = random.randint(min(population), max(population) - 1)
        max_n = random.randint(min_n +  1, max(population))
        
        outcome = [f'between {min_n} and {max_n}, both included', 
                   [x for x in population 
                    if x >= min_n and x <= max_n]]
        outcomes.append(outcome)
        
        outcome = [f'between {min_n} and {max_n}, both excluded', 
                   [x for x in population 
                    if x > min_n and x < max_n]]
        outcomes.append(outcome)
        
        return population, outcomes
        pass
    
    pass
#class ProbabilityIntroHelper:
