import pandas as pd


#GLOBAL_PATH = Path("C:/-----Enter your path to the project folder here-----")


hs4_name = pd.read_excel("raw data/hs4.xlsx", sheet_name= "FAS Value", dtype={"HTS Number":str}, skiprows= 2)
hs4_name = hs4_name.rename(columns={"HTS Number": "hs4"})


#data layouts: https://www.census.gov/foreign-trade/reference/products/layouts/sths6.html
#country codes: https://www.census.gov/foreign-trade/schedules/c/countryname.html
#another source of aggregated exports by state: https://www.census.gov/foreign-trade/statistics/state/origin_movement/index.html
col_num = [
    (0,6),(6,10),(10,12),(12,16),(16,18),
    (18,33),(33,48),(48,63),(63,78),(78,93),
    (93,108),(108,123),(123,138),(138,153),
    (153,168),(168,183),(183,198),(198,213),(213,228)
]

col_name = [
    "hs6","country","state","year","month",
    "value_mo","air_val_mo","air_swt_mo","ves_val_mo","ves_swt_mo",
    "cnt_val_mo","cnt_swt_mo","value_yr","air_val_yr",
    "air_swt_yr","ves_val_yr","ves_swt_yr","cnt_val_yr","cnt_swt_yr"
]

df = pd.read_fwf("raw data/STHS6M2512.TXT", colspecs=col_num, names=col_name, dtype={"hs6":str})
df["hs4"] = df["hs6"].str[:4]



def aggregate_exports(data, group_by=None, select_country=None, varname = "value_yr"):
    if select_country is not None:
        data = data[data["country"].isin(select_country)]

    if not group_by:
        return df["value_yr"].sum()

    # group and aggregate
    return (
        data
        .groupby(group_by, as_index=False)["value_yr"]
        .sum()
        .sort_values(by = "value_yr", ascending=False)
        .rename(columns={"value_yr": varname})
    )


###
def construct_data(df, merge_on, sortby = None, output_filename = None):

    df_ca = aggregate_exports(df, group_by=[merge_on, "country"], select_country=[1220], varname= "expval_ca")
    df_mx = aggregate_exports(df, group_by=[merge_on, "country"], select_country=[2010], varname= "expval_mx")
    df_camx = aggregate_exports(df, group_by=[merge_on], select_country=[2010, 1220], varname= "expval_camx")
    df_world = aggregate_exports(df, group_by=[merge_on], varname= "expval_world")

    combined_file = (df_world
               .merge(df_camx, on=merge_on, how="outer")
               .merge(df_mx[["expval_mx", merge_on]], on=merge_on, how="outer")
               .merge(df_ca[["expval_ca", merge_on]], on=merge_on, how="outer"))
    
    combined_file["camx_p"] = combined_file["expval_camx"]/combined_file["expval_world"]
    combined_file = combined_file.sort_values(by=sortby, ascending=False)
    
    combined_file.to_excel(output_filename, index = False)
    return combined_file


exp_bystate = construct_data(df, merge_on="state", sortby="expval_world", output_filename="processed data/exports_state.xlsx")
exp_hs4 = construct_data(df, merge_on="hs4", sortby="expval_camx", output_filename="processed data/exports_hs4.xlsx")







###### Detailed exports by state
# Exports to Canada and Mexico, by HS
df_ca_hs4_state = aggregate_exports(df, group_by=["hs4", "state"], select_country=[1220], varname= "expval_ca")
df_mx_hs4_state = aggregate_exports(df, group_by=["hs4", "state"], select_country=[2010], varname= "expval_mx")
df_camx_hs4_state = aggregate_exports(df, group_by=["hs4", "state"], select_country=[2010, 1220], varname= "expval_camx")
df_world_hs4_state = aggregate_exports(df, group_by=["hs4", "state"], varname= "expval_world")

exp_hs4_state = (df_world_hs4_state
                 .merge(df_camx_hs4_state, on=["hs4", "state"], how="outer")
                 .merge(df_ca_hs4_state, on=["hs4", "state"], how="outer")
                 .merge(df_mx_hs4_state, on=["hs4", "state"], how="outer")
                )
    
exp_hs4_state["camx_p"] = exp_hs4_state["expval_camx"]/exp_hs4_state["expval_world"]
exp_hs4_state = (exp_hs4_state.merge(hs4_name[["Description", "hs4"]], on = "hs4"))
print(exp_hs4_state)


######
#specific products of potential interests
#aircraft
aircraft = exp_hs4_state[exp_hs4_state["hs4"] == "8800"]
aircraft.to_excel("processed data/aircraft.xlsx")

#electronic integrated circuits
circuits = exp_hs4_state[exp_hs4_state["hs4"] == "8542"]
circuits.to_excel("processed data/electronic_integrated_circuits.xlsx")

#
data_processing_machines_parts = exp_hs4_state[exp_hs4_state["hs4"].isin(["8473", "8471"])]
data_processing_machines_parts.to_excel("processed data/data_processing_machines_parts.xlsx")


with pd.ExcelWriter("processed data/detailed_exports_by_state_hs4.xlsx") as writer:
    for state, data in exp_hs4_state.groupby("state"):
        data.sort_values("expval_camx", ascending=False).to_excel(
            writer, sheet_name = state, index=False
        )