
from copy import deepcopy
from dataclasses import dataclass
from datetime import datetime
import json
import os
from typing import Any, Callable, Optional
import uuid
# from thescientist.analysis import AnalyticalData
from thescientist.dataset import Dataset
from thescientist.metric import Metric, MetricCollector


class ExperimentalGroup:
    
    def __init__(self, name: str, algorithm: Optional[Callable[(...), Any]] = None, metadata: Optional[dict[str, Any]] = None) -> 'ExperimentalGroup':
        if not metadata:
            metadata = dict()
        
        self.name = name
        self.algorithm = algorithm
        self.metadata = {"name": name, **metadata}

    def __repr__(self) -> str:
        return f"ExperimentalGroup(name={self.name}, algorithm={self.algorithm}, metadata={self.metadata})"

    @classmethod
    def load_from(cls, file_path: str, algorithm: Optional[Callable[(...), Any]] = None) -> 'ExperimentalGroup':
        with open(file_path, mode="rb") as fp:
            metadata = json.load(fp)
            return cls(name=metadata.get("name", str(uuid.uuid4())),
                       algorithm=algorithm,
                       metadata=metadata)

    def save(self, path_to_directory: str) -> None:
        
        metadata = self.metadata
        metadata["saved_at"] = datetime.now()

        directory = path_to_directory[:-1] if path_to_directory[:-1] == "/" else path_to_directory
        file_path = f"{directory}/{self.name}.group.sci"

        with open(file_path, mode="w") as fp:
            json.dump(metadata, fp)
    
    def run_with(self, *args, **kwargs) -> Any:
        return self.algorithm(*args, **kwargs)


@dataclass
class ExperimentalResult:
    group: ExperimentalGroup
    metric: Metric
    dataset: Dataset
    
    # def to_analytical_data(self) -> AnalyticalData:
    #     return AnalyticalData(metric_name=self.metric.name,
    #                           group_name=self.group.name,
    #                           data=self.dataset.col(self.metric.name))


class Experiment:

    def __init__(self, 
                 name: str, 
                 groups: list[ExperimentalGroup], 
                 metrics: list[Metric], 
                 experimental_function: Callable[[ExperimentalGroup, MetricCollector], None],
                 n_executions: int = 31, 
                 metadata: Optional[dict[str, Any]] = None) -> None:
        
        if not metadata:
            metadata = dict()
        
        self.name = name
        self.groups = groups
        self.metric_collector = MetricCollector(name=name, metrics=metrics)
        self.experimental_function = experimental_function
        self.n_executions = n_executions
        self.metadata = {"name": name, 
                         "groups": [group.name for group in groups], 
                         "metrics": [metric.as_dict() for metric in self.metric_collector.original_metrics],
                         # "experimental_function": experimental_function, 
                         "n_executions": n_executions,
                         **metadata}

        self.results: list[ExperimentalResult] = []
    
    def run_with(self, independent_variables: Any, *args, **kwargs) -> None:
        group_collector = [(group, deepcopy(self.metric_collector)) for group in self.groups]
        for epoch in range(self.n_executions):
            for group, collector in group_collector:
                
                self.experimental_function(group=group, collector=collector, independent_variables=independent_variables, *args, **kwargs)
        
        for group, collector in group_collector:
            for _, metric in collector.metrics.items():
                self.results.append(ExperimentalResult(group=group, metric=metric["schema"], dataset=metric["dataset"]))
    
    def save(self, path_to_directory: str) -> None:
        dt = str(datetime.now()).replace(" ", "_").replace(":", "-").replace(".", "-").replace("+", "-")

        directory = path_to_directory[:-1] if path_to_directory[:-1] == "/" else path_to_directory
        path = f"{directory}/{self.name}/{dt}"

        os.makedirs(path, exist_ok=True)
        
        for result in self.results:
            result.dataset.to_csv(f"{path}/{result.group.name}_{result.metric.name}.data.csv")

        with open(f"{path}/.metadata.experiment.sci", mode="w") as fp:
            json.dump(self.metadata, fp)


if __name__ == "__main__":

    def experiment_function(group: ExperimentalGroup, collector: MetricCollector, independent_variables: list[Any]) -> None:
        result = group.run_with(independent_variables)
        collector.collect({"abs_value": result})
        collector.collect({"rel_value": result})

    def group1(x: int): return 2*x + 1
    def group2(x: int): return x**2 + x + 1
    def group3(x: int): return 3*x**2 - 5*x
    
    experiment = Experiment(name="Absolute-Value", 
                            groups=[ExperimentalGroup(name="decision-tree", algorithm=group1),
                                    ExperimentalGroup(name="svm", algorithm=group2),
                                    ExperimentalGroup(name="mlp", algorithm=group3)],
                            metrics=[Metric("abs_value", int),
                                     Metric("rel_value", int)],
                            experimental_function=experiment_function,
                            n_executions=3)

    experiment.run_with(independent_variables=2)
    experiment.save("./results")

