import polars as pl
import pandas as pd
from cryptorandom.cryptorandom import SHA256
from scipy.stats import fisher_exact
from scipy.stats.contingency import odds_ratio
import seaborn as sns
from cryptorandom.sample import random_sample
from permute.core import two_sample
import numpy as np

sns.set_theme()

import os
os.chdir('../')
from utils import chi2, permutation

seed = 98066
prng = SHA256(seed)

# data loading
df = pl.read_csv("example-1/example_1_data_sanitized.csv")
shooting_cases = df.filter(pl.col("Someone shot? (0=no, 1=yes)") == 1).filter(pl.col("Exclude?").is_null())#.filter(pl.col("Open  = 0; Closed =1") == 1)

# data loading for permutation test
df_pd = shooting_cases.to_pandas()
df_filtered = df_pd[(df_pd["Someone shot? (0=no, 1=yes)"]==1) & (df_pd["Exclude?"].isnull())].reset_index(drop=True) 
# filter for charge, race
df_filtered = df_filtered[["Docket", "Race", "12022.53(d) alleged"]]
df_white_hispanic = df_filtered[df_filtered["Race"].isin(["W", "H"])].reset_index(drop=True)
df_non_hispanic = df_filtered
df_non_hispanic["race"] = [ "H" if race == "H" else "NH" for race in df_non_hispanic["Race"].to_list()]

shooting_cases_L_W = shooting_cases.group_by(['Race','12022.53(d) alleged'])\
    .len()\
    .filter(pl.col("Race").str.contains_any(["W","H"]))\
    .pivot(["12022.53(d) alleged"], index="Race", values="len")\
    .sort("Race")\
    .rename({"0": "12022.53(d) NOT alleged", "1": "12022.53(d) alleged"})

odds_table = shooting_cases_L_W.select("Race","12022.53(d) alleged", "12022.53(d) NOT alleged")
display("Contingency Table For Odds Ratio", odds_table)
odds_ratio_result = odds_ratio(odds_table.drop("Race"))
print(f"The Odds of {odds_table.columns[1]} for {odds_table[0,"Race"]} is {odds_ratio_result.statistic} times that of {odds_table[1,"Race"]} ")

# risk ratio
t = odds_table.to_pandas().set_index("Race")
(t.at["H", "12022.53(d) alleged"]  / (t.at["H", "12022.53(d) alleged"] + t.at["H", "12022.53(d) NOT alleged"] ))/ \
(t.at["W", "12022.53(d) alleged"]  / (t.at["W", "12022.53(d) alleged"] + t.at["W", "12022.53(d) NOT alleged"] ))

table = shooting_cases_L_W.to_pandas()[["12022.53(d) alleged", "12022.53(d) NOT alleged", "Race"]].set_index("Race").T
table

# fisher
# fishers exact 
# null hypothesis: hispanics and non-hispanics are charged with 664-187 at the same rate
# alternative hypothesis: hispanics are sentenced are charged with 664-187 at a greater rate than non-hispanics
res = fisher_exact(table, alternative='greater')
print(f"p-value: {round(res.pvalue,5)}, test stat: {res.statistic}")

h_enhance, white_enhance = permutation.get_lists(df_white_hispanic, "Race", "H", "W", "12022.53(d) alleged")
_ = permutation.t_test(h_enhance, white_enhance, "greater", prng=prng, title="Proportion with Enhancement")

# # proportion
# def test_stat(df): 
#     return df['12022.53(d) alleged'].mean()
# n_h = df_white_hispanic["Race"].value_counts()["H"]
# obs_stat = test_stat(df_white_hispanic[df_white_hispanic["Race"]=="H"])
# _, _ = permutation.prop_test(n_h, df_white_hispanic, test_stat, obs_stat, prng=prng, plot_title="Latino v White \n Charged with 12022.53(d)", tail="greater")

shooting_cases_L_NL = shooting_cases.group_by(pl.col('Race').is_in(["H"]),'12022.53(d) alleged')\
    .len()\
    .with_columns(
        pl.when(pl.col('Race') == True)
        .then(pl.lit("Latino"))
        .otherwise(pl.lit("Non-Latino")).alias("Race")
    )\
    .pivot(["12022.53(d) alleged"], index="Race", values="len")\
    .sort("Race")\
    .rename({"0": "12022.53(d) NOT alleged", "1": "12022.53(d) alleged"})
    
display(shooting_cases_L_NL)

odds_table = shooting_cases_L_NL.select("Race","12022.53(d) alleged", "12022.53(d) NOT alleged")
display("Contingency Table For Odds Ratio", odds_table)
odds_ratio_result = odds_ratio(odds_table.drop("Race"))
print(f"The Odds of {odds_table.columns[1]} for {odds_table[0,"Race"]} is {odds_ratio_result.statistic} times that of {odds_table[1,"Race"]} ")

# risk ratio
t = odds_table.to_pandas().set_index("Race")
(t.at["Latino", "12022.53(d) alleged"]  / (t.at["Latino", "12022.53(d) alleged"] + t.at["Latino", "12022.53(d) NOT alleged"] ))/ \
(t.at["Non-Latino", "12022.53(d) alleged"]  / (t.at["Non-Latino", "12022.53(d) alleged"] + t.at["Non-Latino", "12022.53(d) NOT alleged"] ))

table = shooting_cases_L_NL.to_pandas()[["12022.53(d) alleged", "12022.53(d) NOT alleged", "Race"]].set_index("Race").T
table

20+41+22+4

# fisher
# fishers exact 
# null hypothesis: hispanics and non-hispanics are charged with 12022.53d at the same rate
# alternative hypothesis: hispanics are sentenced are charged with 12022.53d at a greater rate than non-hispanics
res = fisher_exact(table, alternative='greater')
print(f"p-value: {round(res.pvalue,5)}, test stat: {res.statistic}")

h_enhance, nh_enhance = permutation.get_lists(df_non_hispanic, "race", "H", "NH", "12022.53(d) alleged")
_ = permutation.t_test(h_enhance, nh_enhance, "greater", prng=prng, title="Proportion with Enhancement")