import copy
import random

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import qgrid

from openfisca_reformator.econometrics import gini, draw_ginis
from openfisca_reformator.population import Population
from openfisca_reformator.reform import Reform
from openfisca_reformator.utils import show_histogram, percent_diff, scatter_plot, multi_scatter

qgrid.nbinstall(overwrite=True)

%matplotlib inline

"""
Explanation: Guillaume Chaslot &#103;&#117;&#105;&#108;&#108;&#97;&#117;&#109;&#101;&#46;&#99;&#104;&#97;&#115;&#108;&#111;&#116;&#64;&#100;&#97;&#116;&#97;&#46;&#103;&#111;&#117;&#118;&#46;&#102;&#114;
Tax Legislation Reform
The tax and benefit legislation is complex. This tool helps to make reform that returns a complete and simplified tax and benefit system.
How it works:

Define your concepts (e.g., 'nb_enfants', 'age', 'parent isolé') and budget (e.g.: cost less than 2 millions euros)
The machine learning algorithm helps you adjust the parameters of your reform to approximate the current legislation and fit your budget
From the biggest discrepencies with the current legislation, you can improve your concepts (or the legislation)
Repeat until you reach a legislation that matches your own goals. The algorithm takes care of minimizing your budget and maximizing the similarity to current legislation.

Beta version limitations:
For this test, we only take a population of people from all ages, who have 0 to 5 children, no salary, married, veuf, paces, divorces or celibataires, and simulate the "aides sociales"
Current Results:
Within a few minutes, we got a tax reform for "aides sociales" for people with no salary that is:

6 lines long 
similar to the existing in average at more than 95%

Imports
End of explanation
"""


population = Population.from_openfisca_files(
    base_name='1aj-1bj-f-2000',
    variables_loaded_from_results=['revdisp'],
    nb_individuals=2863.0,
    nb_individuals_ref= 62000000 * 93 / 100,
    )

"""
Explanation: Which population do you want to use?
For instance we want to use the 'revdisp' computed by OpenFisca, on a population that represents 0.21% of the french population.
End of explanation
"""


population = Population.from_csv(
    base_name = 'population_simulated_irpp_fam6310.csv',
    nb_individuals = 6380.0,
    nb_individuals_ref = 62000000 * 93 / 100,
)

"""
Explanation: Or from csv with another simulated population:
End of explanation
"""


def age_dec_1(family):
    return age_from_box(family, '0DA')

def age_dec_2(family):
    if '0DB' in family:
        return age_from_box(family, '0DB')
    return None

def if_married(family):
    return 'M' in family

def if_pacse(family):
    return 'O' in family

def if_veuf(family):
    return 'V' in family

def if_divorce(family):
    return 'D' in family

def if_two_people(family):
    return '0DB' in family

def revenu_1(family):
    return family['1AJ']

def taxable_variable(family):
    if '1BJ' in family:
        return family['1AJ'] + family['1BJ']
    else:
        return family['1AJ']

def revenu_2(family):
    if '1BJ' in family:
        return family['1BJ']
    return False

def age_from_box(family, birthdate_box):
    return 2014 - family.get(birthdate_box, 2014)

def if_both_declarant_parent_below_24(family):
    if age_dec_1(family) >= 24:
        return False
    if '0DB' in family and age_dec_2(family) >= 24:
        return False
    if 'F' not in family or family['F'] == 0:
        return False
    return True

def per_child(family):
    if 'F' in family:
        return family['F']
    else:
        return None

def per_child_parent_isole(family):
    if '0DB' in family:
        return None
    if 'F' in family:
        return family['F']
    else:
        return None

def if_parent_isole_moins_20k(family):
    return 'F' in family and family['F'] >= 1 and ('0DB' not in family) and family['1AJ'] < 20000

def if_enfant_unique(family):
    return 'F' in family and family['F'] == 1

def if_deux_enfants_ou_plus(family):
    return 'F' in family and family['F'] >= 2

def per_child_after_2(family):
    if 'F' in family and family['F'] >= 2:
        return family['F'] - 2
    return None

def if_declarant_above_65(family):
    return age_from_box(family, '0DA') >= 65

def if_codeclarant_above_65(family):
    return '0DB' in family and age_from_box(family, '0DB') > 65

def per_declarant_above_65(family):
    return int(age_from_box(family, '0DB') >= 65) + int(age_from_box(family, '0DA') >= 65)

def per_declarant_above_24(family):
    return int(age_from_box(family, '0DB') >= 24) + int(age_from_box(family, '0DA') >= 24)

def if_one_declarant_above_24(family):
    return (age_from_box(family, '0DA') >= 24 or ('0DB' in family and age_from_box(family, '0DB') >= 24))

def if_one_declarant_above_24_or_has_children(family):
    return (age_from_box(family, '0DA') >= 24 or ('0DB' in family and age_from_box(family, '0DB') >= 24) or 
           ('F' in family and family['F'] >= 1))

def if_earns_10k(family):
    return taxable_variable(family) > 10000

def if_earns_30k(family):
    return taxable_variable(family) > 30000

def if_earns_40k(family):
    return taxable_variable(family) > 40000

def anyone(family):
    return True

population.declare_computed_variables({
    'anyone': anyone,
    'age_dec_1': age_dec_1,
    'age_dec_2': age_dec_2,
    'taxable_variable': taxable_variable,
    'revenu_1': revenu_1,
    'revenu_2': revenu_2,
    'per_child': per_child,
    'per_child_parent_isole': per_child_parent_isole,
    'per_declarant_above_24': per_declarant_above_24,
    'per_declarant_above_65': per_declarant_above_65,
    'per_child_after_2': per_child_after_2,
    'if_earns_10k': if_earns_10k,
    'if_earns_30k': if_earns_30k,
    'if_earns_40k': if_earns_40k,
    'if_one_declarant_above_24': if_one_declarant_above_24,
    'if_one_declarant_above_24_or_has_children': if_one_declarant_above_24_or_has_children,
    'if_both_declarant_parent_below_24': if_both_declarant_parent_below_24,
    'if_declarant_above_65': if_declarant_above_65,
    'if_codeclarant_above_65': if_codeclarant_above_65,
    'if_enfant_unique': if_enfant_unique,
    'if_deux_enfants_ou_plus': if_deux_enfants_ou_plus,
    'if_two_people': if_two_people,
    'if_parent_isole_moins_20k': if_parent_isole_moins_20k,
    })



population.print_population_from_cerfa_values()

population._population

def anyone(family):
    return True

def taxable_variable(family):
    return family['salaire'] + family['salaire_conj']

def taxable_variable(family):
    return family['salaire'] + family['salaire_conj']

def revdisp(family):
    return family['salaire'] + family['salaire_conj'] + family['retraite'] + family['retraite_conj'] + family['irpp']


population.declare_computed_variables({
        'anyone': anyone,
        'taxable_variable': taxable_variable
    })


"""
Explanation: What are the computed variables you want to use for your reform?
You can add concepts like the number of children, age, family situation, etc...
End of explanation
"""


revdisp = list(family['revdisp'] for family in population._population)
show_histogram(revdisp, 'Distribution of revenu disponible')


"""
Explanation: Plots Revenu disponible before reform
End of explanation
"""


reform = Reform(
    population=population,
    target_variable='revdisp',
    taxable_variable='taxable_variable',

    # Constraints
    max_cost=0,
    benefits_constraints=[
        'per_child',
        'per_declarant_above_65',
        'if_one_declarant_above_24_or_has_children',
        ],
    taxes_constraints=[
        'anyone',
        'per_child',
        ],
    tax_threshold_constraints=[
        'anyone',
        'per_child'
        ],
    max_evals=5000
    )

reform.show_stats()

reform.draw_ginis()

reform.show_constraints_value('benefit')

reform.show_constraints_value('tax_rate')

reform.show_constraints_value('tax_threshold')

new_pop = copy.deepcopy(population._population)

for i in range(len(new_pop)):
    new_pop[i]['reform'] = reform.simulated_results[i] 

def plot_for_population(pop):
    revenu_imposable = list(family['taxable_variable'] for family in pop)
    revdisp = list(family['revdisp'] for family in pop)
    reform = list(family['reform'] for family in pop)

    multi_scatter('Revenu disponible for different people before / after reforme', 'Revenu initial', 'Revenu disponible', [
                  {'x':revenu_imposable, 'y':reform, 'label':'Reform', 'color':'blue'},
                  {'x':revenu_imposable, 'y':revdisp, 'label':'Original', 'color':'red'},
              ])

un_enfant_pop = list(filter(lambda x: x.get('per_child', 0) == 0
                                  and x.get('age_dec_1', 0) < 65
                                  and x.get('age_dec_2', 0) < 65,
                            new_pop))

plot_for_population(un_enfant_pop)

"""
Explanation: <font color='blue'>Define the constraints of your reform</font>
End of explanation
"""


xmin = 0
xmax = 60000
nb_buckets = 35

bins = np.linspace(xmin, xmax, nb_buckets)

plt.hist(revdisp, bins, alpha=0.5, label='current')
plt.hist(reform.simulated_results, bins, alpha=0.5, label='reform')
plt.legend(loc='upper right')
plt.show()

"""
Explanation: Plots revenu disponible after reform
End of explanation
"""


difference = list(reform.simulated_results[i] - revdisp[i] for i in range(len(reform.simulated_results)))

show_histogram(difference, 'Changes in revenu')

"""
Explanation: Distribution of the changes in revenu in euros
End of explanation
"""


percentage_difference = list(100 * percent_diff(reform.simulated_results[i], revdisp[i]) for i in range(len(reform.simulated_results)))

show_histogram(percentage_difference, 'Changes in revenu')

"""
Explanation: Distribution of the change in revenu in percentage
End of explanation
"""


nb_children = list((population._population[i].get('per_child', 0)  for i in range(len(population._population)))) 

scatter_plot(nb_children, difference, 'Children', 'Difference reform - current', alpha=0.05)

"""
Explanation: Change as a function of the number of children
End of explanation
"""


age_dec1 = list((population._population[i].get('age_dec_1', 0)  for i in range(len(population._population)))) 

scatter_plot(age_dec1, difference, 'Age declarant 1', 'Difference reform - current', alpha=0.1)

"""
Explanation: Change as a function of the age of declarant 1
A scatter plot is better than a thousand points #ChineseProverb
End of explanation
"""



order = sorted(range(len(reform.simulated_results)), key=lambda k: -(reform.simulated_results[k] - revdisp[k]))

data = {}

possible_keys = set()
for i in order:
    for key in population._raw_population[i]:
        if key != 'year':
            possible_keys.add(key)

for i in order:
    # Adding the diff with the reform.
    differences = data.get('difference', [])
    differences.append(int(reform.simulated_results[i] - revdisp[i]))
    data['difference'] = differences

    for key in possible_keys:
        new_vals = data.get(key, [])
        value = population._raw_population[i].get(key, '')
        if type(value) == float:
            value = int(value)
        new_vals.append(value)
        data[key] = new_vals

    # Adding reformed line.
    reforms = data.get('reform', [])
    reforms.append(int(reform.simulated_results[i]))
    data['reform'] = reforms

df = pd.DataFrame(data=data)
df.set_index('difference', inplace=True)
qgrid.show_grid(df)


"""
Explanation: <font color="blue">Most important: Edge Cases</font>
This is the heart of this tool: by seeing the worse cases, you can discover when the current legislation is smarter than yours (or the other way), and improve it further
End of explanation
"""


