import os
os.chdir("../../")

import pandas as pd
import numpy as np
import statsmodels.formula.api as smf
from scripts.python.tsa.mtsmodel import *
from scripts.python.tsa.ts_eval import *

import seaborn as sns
sns.set_style("whitegrid")
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings("ignore")

for country in ["palau", "samoa", "tonga", "solomon_islands", "vanuatu"]:

    folderpath = os.getcwd() + "/data/tourism/" + str(country) + "/model/"
    mappings = [("var", "pred_total"),
                ("sarimax", "train_pred"),
                ("lf", "pred_mean")]

    country_pred = pd.DataFrame()
    for mapping in mappings:
        model, column = mapping
        filepath = folderpath + str(model) + "_" + str(country) + ".csv"
        pred_df = pd.read_csv(filepath).drop("Unnamed: 0", axis=1)
        pred_df["date"] = pd.to_datetime(pred_df["date"])

        model_col = (pred_df[["date", "total", column]]
                     .rename({column: model}, axis=1))

        if country_pred.empty:
            country_pred = model_col
        else:
            country_pred = country_pred.merge(
                model_col, how="left", on=["date", "total"]).fillna(0)

    # Mean
    country_pred["mean_ensemble"] = (
        country_pred[["sarimax", "var", "lf"]].mean(axis=1))

    # Relative Performance Weights
    country_pred["rpw"], weights = get_rpw(country_pred)

    # OLS (regularized) without intercept
    ols = smf.ols("total~sarimax+var+lf - 1", data=country_pred)
    ols_res = ols.fit()
    ols_lasso = ols.fit_regularized()
    country_pred["ols"] = ols_res.fittedvalues
    country_pred["ols_lasso"] = ols_lasso.fittedvalues

    methods = ["sarimax", "var", "lf"]
    constrained_ls, cls_pred = get_constrained_ls(y=country_pred["total"],
                                                  X=country_pred[methods])
    country_pred["cls"] = cls_pred

    for method, cls in zip(methods, constrained_ls):
        weights["ols_"+str(method)] = ols_res.params[method]
        weights["ols_lasso_"+str(method)] = ols_lasso.params[method]
        weights["cls_"+str(method)] = cls

    weights["date"] = country_pred["date"]

    # Sort columns
    cols = weights.columns.tolist()
    cols = cols[-1:] + cols[:-1]
    weights = weights[cols]

    # Save Combination weights
    weights.to_csv(folderpath+"forecast_combo_weights.csv",
                   encoding="utf-8")

    # Save Forecast Prediction
    country_pred.to_csv(folderpath+"forecast_combo.csv",
                        encoding="utf-8")

    evals = pd.DataFrame()
    for col in ["sarimax", "var", "lf", "mean_ensemble", "rpw", "ols", "cls", "ols_lasso"]:
        mod_eval = pd.DataFrame(calculate_evaluation(country_pred["total"], country_pred[col]),
                                index=[col])
        evals = pd.concat([evals, mod_eval], axis=0)

    evals.columns.name = str(country)
    evals = evals.style.apply(
        lambda x: ['background-color: yellow' if v == x.min() else '' for v in x])
    display(evals)