# Import Libraries
import microdf as mdf
import numpy as np
import pandas as pd
import plotly.express as px

# Import data
raw = pd.read_csv("https://github.com/UBICenter/Veteran-s_Day/raw/main/VeteranData.gz")

# Create Demographic Columns
person = raw.copy(deep=True)
person.columns = person.columns.str.lower()
person["child"] = person.age < 18
person["adult"] = person.age >= 18
person["veteran"] = (person.vetstat == 2) & person.adult
person["non_veteran"] = (person.vetstat == 1) & person.adult

veterans = person.groupby(["spmfamunit"])[["veteran"]].sum()
veterans.columns = ["total_veterans"]
person = person.merge(veterans.reset_index(), on=["spmfamunit"])

person["child_with_vet"] = (person.child) & (person.total_veterans > 0)
person["child_with_no_vet"] = (person.child) & (person.total_veterans == 0)

# Show total veterans in millions
total_veterans = (person.veteran * person.asecwt).sum()
total_veterans / 1_000_000

# Calculate total AGI
person["adjginc"].replace({99999999: 0}, inplace=True)
population = person.asecwt.sum()
person["weighted_agi"] = person.adjginc * person.asecwt
total_agi = person.weighted_agi.sum()

# Calculate AGI tax rate per dollar of UBI
fed_tax_rate_per_dollar_ubi_monthly = (population * 12) / total_agi

# Create table showing tax amounts
tax_rates = pd.DataFrame(np.arange(0, 1001, 50))
tax_rates.columns = ["monthly_ubi"]


def tax(monthly_ubi):
    return (monthly_ubi * fed_tax_rate_per_dollar_ubi_monthly * 100).round(1)


def tax_row(row):
    return tax(row.monthly_ubi)


tax_rates["tax_rate"] = tax_rates.apply(tax_row, axis=1)
tax_rates.columns = ["Monthly UBI", "Flat Tax Rate on AGI"]


def ubi(status, monthly_ubi):
    """At a given UBI level, calculate the poverty rate, median resources,
    mean resources, and percent of people better off for:
    * Veterans
    * Non-Veterans
    * Children living with Veteran
    * Children not living with Veterans
  
    Args:
        status: A person's Veteran status.
            For this simulation their are 4 categories:
            * Veteran
            * Non-Veteran
            * Child living with a Veteran
            * Children not living with a Veteran
        monthly_ubi: the monthly cash transfer given to each person
  
    Returns:
        pandas Series with four elements for the selected group:
        * Poverty rate
        * Median resources per person
        * Mean resources per person
        * Percent of people better off
    """

    # Create a copy of the person DataFrame
    target_persons = person.copy(deep=True)

    # Calculate a person's tax increase
    target_persons["tax_increase"] = (
        fed_tax_rate_per_dollar_ubi_monthly * monthly_ubi * target_persons.adjginc
    )

    # Calculate the total UBI per SPM unit.
    target_persons["total_ubi"] = target_persons.spmnpers * 12 * monthly_ubi

    # Calculate the total tax increase of an SPM unit
    spmu = target_persons.groupby(["spmfamunit"])[["tax_increase"]].sum()
    spmu.columns = ["total_tax_increase"]
    target_persons = target_persons.merge(
        spmu, left_on=["spmfamunit"], right_index=True
    )

    # Calculate each SPM unit's tax rate person
    target_persons["new_spm_resources"] = (
        target_persons.spmtotres
        + target_persons.total_ubi
        - target_persons.total_tax_increase
    )

    # Calculate the new resources per person of each SPM unit
    target_persons["new_resources_per_person"] = (
        target_persons.new_spm_resources / target_persons.spmnpers
    )

    # Slice the data based on Race input
    if status == "veteran":
        target_persons = target_persons[target_persons.veteran]
    if status == "non_veteran":
        target_persons = target_persons[target_persons.non_veteran]

    if status == "veteran_child":
        target_persons = target_persons[target_persons.child_with_vet]
    if status == "non_veteran_child":
        target_persons = target_persons[target_persons.child_with_no_vet]

    # Calculate the change in poverty rate
    target_persons["poor"] = target_persons.new_spm_resources < target_persons.spmthresh
    total_poor = (target_persons.poor * target_persons.asecwt).sum()
    target_pop = target_persons.asecwt.sum()

    # Calculate percent better off
    target_persons["better_off"] = (
        target_persons.new_spm_resources > target_persons.spmtotres
    )
    total_better_off = (target_persons.better_off * target_persons.asecwt).sum()
    percent = total_better_off / target_pop * 100

    return pd.Series(
        [
            mdf.weighted_median(
                target_persons, "new_resources_per_person", "asecwt"
            ).round(0),
            mdf.weighted_mean(
                target_persons, "new_resources_per_person", "asecwt"
            ).round(0),
            (total_poor / target_pop * 100).round(1),
            percent,
        ]
    )


def ubi_row(row):
    """ run the ubi function across the rows of a DataFrame.
  
    Arguments
    row: the row of the DataFrame containing a person's race and the monthly UBI amount
  
    Returns
    The poverty rate for the selected row.
    The median resources per person for the selected row.
    The mean resources per person for the selected row.
    The percent of people better off under the program for the selected row.
    """
    return ubi(row.status, row.monthly_ubi)


# Create a DataFrame that has each the each monthly UBI amount for each race input
summary = mdf.cartesian_product(
    {
        "monthly_ubi": np.arange(0, 1001, 50),
        "status": ["veteran", "non_veteran", "veteran_child", "non_veteran_child"],
    }
)

# Calculate the poverty rate for each row of the summary DataFrame
summary[
    [
        "med_resources_per_person",
        "mean_resources_per_person",
        "poverty_rate",
        "better_off",
    ]
] = summary.apply(ubi_row, axis=1)

summary2 = summary[summary["status"] != "non_veteran_child"]
summary2 = summary2[summary2["status"] != "veteran_child"]
# Format text
center = {
    "med_resources_per_person": "Median resources",
    "mean_resources_per_person": "Mean resources",
}
status = {
    "veteran": "Veterans",
    "non_veteran": "Non-Veteran adults",
    "veteran_child": "Children in Veteran households",
    "non_veteran_child": "Children in non-Veteran households",
}

summary["status"] = summary["status"].map(status)
summary2["status"] = summary2["status"].map(status)

# Colors from https://material.io/design/color/the-color-system.html
BLUE = "#1976D2"
DARK_BLUE = "#0D47A1"
BARELY_BLUE = "#64B5F6"
GRAY = "#9E9E9E"
LIGHT_GRAY = "#E0E0E0"

COLOR_MAP = {
    "Veterans": BLUE,
    "Non-Veteran adults": GRAY,
    "Children in Veteran households": BARELY_BLUE,
    "Children in non-Veteran households": LIGHT_GRAY,
}


def line_graph(df, x, y, color, title, xaxis_title, yaxis_title):
    """Style for line graphs.
    
    Args:
        df: DataFrame with data to be plotted.
        x: The string representing the column in df that holds the new
            spending in billions.
        y: The string representing the column in df that holds the poverty
            rate.
        color: The string representing the UBI type.
        xaxis_title: The string represnting the xaxis-title.
        yaxis_title: The string representing the yaxis-title.
    
    Returns:
        Nothing. Shows the plot.
    """
    fig = px.line(df, x=x, y=y, color=color, color_discrete_map=COLOR_MAP)
    fig.update_layout(
        title=title,
        xaxis_title=xaxis_title,
        yaxis_title=yaxis_title,
        yaxis_ticksuffix="%",
        font=dict(family="Roboto"),
        hovermode="x",
        xaxis_tickprefix="$",
        xaxis_ticksuffix="",
        plot_bgcolor="white",
        legend_title_text="",
    )
    fig.update_traces(mode="markers+lines", hovertemplate=None)

    hide_line = ["Children in Veteran households", "Children in non-Veteran households"]

    fig.for_each_trace(
        lambda trace: trace.update(visible="legendonly")
        if trace.name in hide_line
        else ()
    )

    return fig


fig = line_graph(
    df=summary,
    x="monthly_ubi",
    y="poverty_rate",
    color="status",
    title="The impact of a UBI on Veterans and their families",
    xaxis_title="Monthly UBI",
    yaxis_title="SPM poverty rate",
)
fig.show()