import unittest

import numpy as np
import pandas as pd
from pathlib import Path
from numpy.linalg import lstsq
import json

from skmixed.lme.problems import LinearLMEProblem
from skmixed.lme.cascades import LMECascade


class TestLMECascade(unittest.TestCase):

    def test_solving_marlena_dataset(self):
        # This dataset was generated by this generator:
        # https://github.com/ihmeuw-msca/multilevel-model
        # The whole class appeared as a solution for a
        # small internal Math Sciences team competition in Sept. 2020.

        # Paths to dataset and true parameters/answers
        dataset = "simulation-2"
        dataset_generator_directory = Path("examples/cascades/datasets")
        dataset_path = dataset_generator_directory / dataset / (dataset + ".csv")
        hierarchy_path = dataset_generator_directory / "hierarchy.csv"
        true_coefficients_path = dataset_generator_directory / dataset / "KEY" / ("PARAMS-" + dataset + ".csv")
        true_answers_path = dataset_generator_directory / dataset / "KEY" / ("TRUTH-" + dataset + ".csv")
        true_parameters_path = dataset_generator_directory / dataset / ("SETTINGS-" + dataset + ".json")

        # Reading data
        data = pd.read_csv(dataset_path)
        hierarchy = pd.read_csv(hierarchy_path)
        hierarchy = hierarchy[hierarchy.location_id.isin(data.location_id.unique())]
        hierarchy = hierarchy[['location_name', 'location_id', 'parent_id', 'level']].reset_index(drop=True)
        data = data.merge(hierarchy, on="location_id", how="left")

        levels = sorted(data["level"].unique())

        covariates_colnames = [column for column in data.columns if column.startswith("x_")]

        data["intercept"] = 1
        data["predictions"] = 0
        data["residues"] = 0
        data["se"] = 0
        hierarchy["se"] = 1

        # Transform hierarchy into dict: {child_id: parent_id}
        hierarchy_dict = {row["location_id"]: row["parent_id"] for i, row in hierarchy.iterrows()}

        # Inferring observational standard error
        # NB: it's completely unnecessary to do it in this way:
        # you can (and should) estimate it along with beta and gamma.
        # However, this module was a hackaton-style work, so I did not have time
        # to implement it.
        current_level = data["level"] == 0
        level0_data = data[current_level][["intercept"] + covariates_colnames].to_numpy()
        level0_target = data[current_level]["observation"].to_numpy()
        level0_coefficients = lstsq(level0_data, level0_target, rcond=None)[0]
        data.loc[current_level, "predictions"] = level0_data.dot(level0_coefficients)
        data.loc[current_level, "residues"] = data.loc[current_level, "observation"] - data.loc[
            current_level, "predictions"]
        current_se = np.sqrt((data.loc[current_level, "residues"].var()))

        # Generating one problem per level
        problems = []
        for level in levels:
            current_level = data["level"] == level
            data.loc[current_level, "se"] = current_se
            X = data[current_level][covariates_colnames + ["location_id", "se"]].to_numpy()
            column_labels = [1] * len(covariates_colnames) + [0, 4]
            X = np.vstack([column_labels, X])
            y = data[current_level]["observation"].to_numpy()
            problem = LinearLMEProblem.from_x_y(X, y, random_intercept=True)
            problems.append(problem)

        # Fitting a model
        model = LMECascade()
        model.fit_problem(problems, hierarchy_dict)
        beta = model.coef_["beta"]

        # Aggregating answers into a dataframe
        answers = [{"location_id": group, "intercept": beta[0] + u[0], "effect_0": beta[1]} for group, u in
                   model.coef_["random_effects"].items()]
        coefficients = pd.DataFrame(answers)
        true_coefficients = pd.read_csv(true_coefficients_path)
        coefficients = coefficients.merge(true_coefficients, on="location_id", how="inner", suffixes=('_pred', '_true'))
        with open(true_parameters_path) as f:
            true_parameters = json.load(f)

        # In case you want to see the results
        # print(coefficients[["location_id", "intercept_true", "intercept_pred", "effect_0_true", "effect_0_pred"]])

        # Checking that the scores on the similation-2 dataset are still
        # at least as good as they were when I finished the contest
        self.assertTrue(np.linalg.norm(coefficients['intercept_true'] - coefficients['intercept_pred']) < 1.62)
        self.assertTrue(np.linalg.norm(coefficients['effect_0_true'] - coefficients['effect_0_pred']) < 0.72)