from darts import TimeSeries
from darts.datasets import ILINetDataset
from darts.metrics import mape
from darts.models import ExponentialSmoothing
from darts.utils.missing_values import fill_missing_values
from darts.dataprocessing.transformers import Scaler
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import os

def load_ILINetDataset():
    """
    Dataset's Components Descriptions:

    * % WEIGHTED ILI: Combined state-specific data of patients visit to healthcare providers for ILI reported each week weighted by state population

    * % UNWEIGHTED ILI: Combined state-specific data of patients visit to healthcare providers for ILI reported each week unweighted by state population

    * AGE 0-4: Number of patients between 0 and 4 years of age

    * AGE 25-49: Number of patients between 25 and 49 years of age

    * AGE 25-64: Number of patients between 25 and 64 years of age

    * AGE 5-24: Number of patients between 5 and 24 years of age

    * AGE 50-64: Number of patients between 50 and 64 years of age

    * AGE 65: Number of patients above (>=65) 65 years of age

    * ILITOTAL: Total number of ILI patients. For this system, ILI is defined as fever (temperature of 100°F [37.8°C] or greater) and a cough and/or a sore throat

    * NUM. OF PROVIDERS: Number of outpatient healthcare providers

    * TOTAL PATIENTS: Total number of patients
    """
    ilidata = ILINetDataset().load().astype(np.float32)

    return ilidata


def preprocess_data(ilidata):
    ilitotal = ilidata['ILITOTAL']
    covariates = ilidata.drop_columns(col_names='ILITOTAL')  # TODO perform predictions with covariates time series

    # Apply logaritmic transformation
    log_ili = ilitotal.map(np.log)

    # Remove -inf values
    pd_log_ili = log_ili.pd_dataframe()
    pd_log_ili['ILITOTAL'].replace(to_replace=-np.inf, value=np.nan, inplace=True)
    log_ili = TimeSeries.from_dataframe(pd_log_ili)

    log_ili = fill_missing_values(log_ili)

    return log_ili


def train_val_split(ilidata):
    train_ili, val_ili = ilidata.split_before(0.75)

    return train_ili, val_ili


def scale_train(train_ili):
    scaler = Scaler()
    train_ili_scaled = scaler.fit_transform(train_ili)

    return train_ili_scaled, scaler


def train(train_ili_scaled):
    model = ExponentialSmoothing()
    model.fit(train_ili_scaled)

    return model


def save_model(model, path):
    model_name = str(model).split('(')[0]
    if not os.path.exists(path=path):
        os.makedirs(path)
        model.save(os.path.join(path, model_name))
    #else:
    #    model = ExponentialSmoothing.load(path)


def load_model(path):
    model = ExponentialSmoothing.load(path)
    
    return model


def predict(model, val_ili):
    preds = model.predict(len(val_ili))

    return preds


def inverse_scale_predictions(scaled_preds, scaler):
    preds = scaler.inverse_transform(scaled_preds)

    return preds


def plot_results(train_ili, val_ili, preds):
    fig, _ = plt.subplots()
    train_ili.plot(label='train')
    val_ili.plot(label='val')
    preds.plot(label='preds')
    
    return fig


def compute_mape(preds, val):
    metric = mape(val, preds)

    return metric