import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

from mlc.constraints.relation_constraint import Constant, Feature, SafeDivision
from mlc.datasets.dataset import (
    CsvDataSource,
    Dataset,
    DateSorter,
    DownloadFileDataSource,
    Splitter,
    Task,
)


class LcldSplitterTime(Splitter):
    def get_splits(self, dataset):
        _, _, t = dataset.get_x_y_t()
        i_train = np.where(
            (pd.to_datetime("2013-01-01") <= t)
            & (t <= pd.to_datetime("2015-06-30"))
        )[0]
        i_val = np.where(
            (pd.to_datetime("2015-07-01") <= t)
            & (t <= pd.to_datetime("2015-12-31"))
        )[0]
        i_test = np.where(
            (pd.to_datetime("2016-01-01") <= t)
            & (t <= pd.to_datetime("2017-12-31"))
        )[0]
        return {"train": i_train, "val": i_val, "test": i_test}


class LcldSplitterIid(LcldSplitterTime):
    def get_splits(self, dataset):
        _, y, _ = dataset.get_x_y_t()
        splits = super(LcldSplitterIid, self).get_splits(dataset)
        merged = np.concatenate(
            [splits.get("train"), splits.get("val"), splits.get("test")]
        )
        i_train, i_test = train_test_split(
            merged,
            random_state=100,
            shuffle=True,
            stratify=y[merged],
            test_size=splits.get("test").shape[0],
        )
        i_train, i_val = train_test_split(
            i_train,
            random_state=200,
            shuffle=True,
            stratify=y[i_train],
            test_size=splits.get("val").shape[0],
        )
        return {"train": i_train, "val": i_val, "test": i_test}


def get_relation_constraints():

    int_rate = Feature("int_rate") / Constant(1200)
    term = Feature("term")
    installment = Feature("loan_amnt") * (
        (int_rate * ((Constant(1) + int_rate) ** term))
        / ((Constant(1) + int_rate) ** term - Constant(1))
    )
    g1 = Feature("installment") == installment

    g2 = Feature("open_acc") <= Feature("total_acc")

    g3 = Feature("pub_rec_bankruptcies") <= Feature("pub_rec")

    g4 = (Feature("term") == Constant(36)) | (Feature("term") == Constant(60))

    g5 = Feature("ratio_loan_amnt_annual_inc") == (
        Feature("loan_amnt") / Feature("annual_inc")
    )

    g6 = Feature("ratio_open_acc_total_acc") == (
        Feature("open_acc") / Feature("total_acc")
    )

    # g7 was diff_issue_d_earliest_cr_line
    # g7 is not necessary in v2
    # issue_d and d_earliest cr_line are replaced
    # by month_since_earliest_cr_line

    g8 = Feature("ratio_pub_rec_month_since_earliest_cr_line") == (
        Feature("pub_rec") / Feature("month_since_earliest_cr_line")
    )

    g9 = Feature(
        "ratio_pub_rec_bankruptcies_month_since_earliest_cr_line"
    ) == (
        Feature("pub_rec_bankruptcies")
        / Feature("month_since_earliest_cr_line")
    )

    g10 = Feature("ratio_pub_rec_bankruptcies_pub_rec") == SafeDivision(
        Feature("pub_rec_bankruptcies"), Feature("pub_rec"), Constant(-1)
    )

    return [g1, g2, g3, g4, g5, g6, g8, g9, g10]


data_source = DownloadFileDataSource(
    url="https://uniluxembourg-my.sharepoint.com/:x:/g/personal/"
    "thibault_simonetto_uni_lu/"
    "EeKpMXnQNo9CuRaLcHNjiX0B4Tf2H_HV3OKmlqvwbZZ-aA?download=1",
    file_data_source=CsvDataSource(path="./data/mlc/lcld_v2/lcld_v2.csv"),
)
metadata_source = DownloadFileDataSource(
    url="https://uniluxembourg-my.sharepoint.com/:x:/g/personal/"
    "thibault_simonetto_uni_lu/"
    "EfIRH-UBrW1BiYwvSjlyZIIBFbOiALcxUzdZ5T11qhILlw?download=1",
    file_data_source=CsvDataSource(
        path="./data/mlc/lcld_v2/lcld_v2_metadata.csv"
    ),
)
tasks = [
    Task(
        name="charged_off",
        task_type="classification",
        evaluation_metric="f1_score",
    )
]
sorter = DateSorter(date_col="issue_d")
relation_constraints = get_relation_constraints()


lcld_time = Dataset(
    name="lcld_v2_time",
    data_source=data_source,
    metadata_source=metadata_source,
    tasks=tasks,
    sorter=sorter,
    splitter=LcldSplitterTime(),
    relation_constraints=get_relation_constraints(),
)

lcld_iid = Dataset(
    name="lcld_v2_iid",
    data_source=data_source,
    metadata_source=metadata_source,
    tasks=tasks,
    sorter=sorter,
    splitter=LcldSplitterIid(),
    relation_constraints=get_relation_constraints(),
)

datasets = [lcld_time, lcld_iid]
