import os, shutil
import os.path as osp
import pickle
import numpy as np
from collections import deque
from dataclasses import dataclass, field
from functools import cached_property
from pydantic import BaseModel
from typing import TypeVar, Generic, Dict, Optional
from scipy.stats import ttest_rel
from copy import deepcopy
from abc4d.core.dataloader import StepBasedDataLoader
from abc4d.utils.plog import ProgressLogger
from abc4d.utils.opt_chain import OptChain


T = TypeVar('T')


class AB(Generic[T]):
    """
    The container to specify components for A/B testing.
    where A is the control group and B is the test group.
    """

    def __init__(self, *args: T):
        if len(args) == 2:
            self.A = args[0]
            self.B = args[1]
        elif len(args) == 1:
            self.A = args[0]
            self.B = args[0]
        else:
            raise ValueError("Invalid number of arguments")


class ModelState(BaseModel):
    global_step: int
    trained_step: int
    # global_step -> metric_name -> MetricSpec
    history: Dict[int, Dict[str, "MetricSpec"]]


class Model:

    _summary = ""

    def __init__(self, *, path: str = "", migrate: bool = True):
        self.state = ModelState(
            global_step=0, trained_step=0, history={})
        self.build_model()
        self.build_optimizer()
        if path:
            self.load(path, migrate=migrate)
    
    def build_model(self):
        raise NotImplementedError()
        
    def build_optimizer(self):
        """build optimizer for train from scratch"""
        raise NotImplementedError()
    
    def migrate_baseline_params(self, checkpoint):
        return checkpoint
        
    def load(self, path: str, migrate: bool = True):
        path = path if path.endswith(".pkl") else path + ".pkl"
        with open(path, "rb") as f:
            full_state = pickle.load(f)
        self.state = ModelState.model_validate(full_state["state"])
        if migrate:
            self.load_checkpoint(self.migrate_baseline_params(full_state["checkpoint"]))
        else:
            self.load_checkpoint(full_state["checkpoint"])
        return self

    def save(self, path: str, global_step: int):
        save_path = path + f"_{global_step}.pkl"
        symlink_path = path + ".pkl"
        full_state = {"state": self.state.model_dump(), "checkpoint": self.save_checkpoint()}
        with open(save_path, "wb") as f:
            pickle.dump(full_state, f)
        if os.path.isfile(symlink_path):
            os.remove(symlink_path)
        os.symlink(save_path, symlink_path)

    def load_checkpoint(self, checkpoint):
        raise NotImplementedError()

    def save_checkpoint(self):
        raise NotImplementedError()

    def freeze_baseline(self):
        raise NotImplementedError()

    def unfreeze_baseline(self):
        raise NotImplementedError()

    def train_on_batch(self, batch) -> Dict[str, "MetricSpec"]:
        raise NotImplementedError()
    
    @property
    def summary(self) -> str:
        return self._summary or self.__class__.__name__
        
    @summary.setter
    def summary(self, value: str):
        self._summary = value


@dataclass
class MetricSpec:
    value: float
    importance: float
    higher_is_better: bool
    min_value: float = 0.
    max_value: float = 1.

    def __post_init__(self):
        if self.min_value is None:
            self.min_value = self.value
        if self.max_value is None:
            self.max_value = self.value

    def update_min_max(self, another: "MetricSpec"):
        self.min_value = min(self.min_value, another.min_value)
        self.max_value = max(self.max_value, another.max_value)
        another.min_value = self.min_value
        another.max_value = self.max_value


@dataclass
class Compare:
    p_value: float
    avg_A: float = 0.5
    avg_B: float = 0.5

    def smoother(self, smoothed_report: "Compare", smooth_factor: float = 0.9):
        # reset p-value if the conclusion is switching
        if (smoothed_report.avg_A - smoothed_report.avg_B) * (self.avg_A - self.avg_B) < 0:
            smoothed_report.p_value = 1.0
            self.p_value = 1.0
    
        return Compare(
            p_value=smoothed_report.p_value * smooth_factor +
            self.p_value * (1 - smooth_factor),
            avg_A=smoothed_report.avg_A * smooth_factor +
            self.avg_A * (1 - smooth_factor),
            avg_B=smoothed_report.avg_B * smooth_factor +
            self.avg_B * (1 - smooth_factor)
        )

    def A_better(self, p_thres: float = 0.05) -> bool:
        return self.avg_A > self.avg_B and self.p_value < p_thres

    def B_better(self, p_thres: float = 0.05) -> bool:
        return self.avg_B > self.avg_A and self.p_value < p_thres

    def no_difference(self, p_thres: float = 0.05) -> bool:
        return self.p_value > p_thres

    def format(self):
        compare = "A is better" if self.A_better(
        ) else "B is better" if self.B_better() else "no difference"
        
        if self.avg_B > self.avg_A:
            return (
            "A: ",
            (f"{self.avg_A:.4f}", "red"),
            ", B: ",
            (f"{self.avg_B:.4f}", "green"),
            ", p-value: ",
            f"{self.p_value:.4f}",
            ", compare: ",
            compare
            )
        else:
            return (
            "A: ",
            (f"{self.avg_A:.4f}", "green"),
            ", B: ",
            (f"{self.avg_B:.4f}", "red"),
            ", p-value: ",
            f"{self.p_value:.4f}",
            ", compare: ",
            compare
            )


class MetricsTracker:

    def __init__(self, window_size: int = 100):
        self.window_size = window_size
        self.buffer = AB(deque(maxlen=window_size), deque(maxlen=window_size))

    def reduce(self, metrics: Dict[str, MetricSpec]) -> float:
        """ weighed sum of normalized metrics """
        weighted_sum = 0.0
        total_importance = 0.0

        for metric in metrics.values():
            normalized_value = (metric.value - metric.min_value) / \
                (metric.max_value - metric.min_value)
            normalized_value = max(min(normalized_value, 1.0), 0.0)
            if not metric.higher_is_better:
                normalized_value = 1 - normalized_value
            weighted_sum += normalized_value * metric.importance
            total_importance += metric.importance

        return weighted_sum / total_importance if total_importance != 0 else 0.0

    def update(self, metrics_A: Dict[str, MetricSpec], metrics_B: Dict[str, MetricSpec]):
        self.buffer.A.append(self.reduce(metrics_A))
        self.buffer.B.append(self.reduce(metrics_B))

    def compare(self) -> Optional[Compare]:
        if len(self.buffer.A) < self.window_size:
            return

        t_stat, p_value = ttest_rel(self.buffer.A, self.buffer.B, nan_policy='omit')
        
        if np.isnan(p_value):
            p_value = 0.5
        
        avg_A = sum(self.buffer.A) / \
            len(self.buffer.A) if self.buffer.A else 0.0
        avg_B = sum(self.buffer.B) / \
            len(self.buffer.B) if self.buffer.B else 0.0

        return Compare(p_value=p_value, avg_A=avg_A, avg_B=avg_B)


class Monitor:

    def log(self, metrics: Dict[str, MetricSpec]):
        pass


class SimpleTrainer:
    
    def __init__(
        self,
        model: Model,
        save_path: str,
        report_interval: int = 100,
    ):
        self.model = model
        self.save_path = save_path
        self.report_interval = report_interval
        
        os.makedirs(self.trial_save_path)
        self.plog = ProgressLogger(osp.join(self.trial_save_path, "progress.log"))
        
    @cached_property
    def trial_save_path(self) -> str:
        # in the save_path, there will be subfolders for each run(trial) with autoincrementing index.
        if osp.exists(self.save_path):
            last_index = max([int(trial_folder[6:]) for trial_folder in os.listdir(self.save_path) if trial_folder.startswith("trial_") and trial_folder[6:].isdigit()])
        else:
            last_index = 0
        return osp.join(self.save_path, f"trial_{last_index+1:02d}")
        
    @property
    def global_step(self):
        return self.model.state.global_step
    
    @global_step.setter
    def global_step(self, value):
        self.model.state.global_step = value

    def train(self, data: StepBasedDataLoader, steps: int):
        plog = self.plog
        plog.info(f"Model: {self.model.summary}")
        
        with plog.progress_bar(range(steps), desc="Training") as pbar:
            for current_step in pbar:
                batch = data.get_batch_at_step(self.global_step)
                metrics = self.model.train_on_batch(batch)
                self.model.state.history[self.global_step] = deepcopy(metrics)
                self.global_step += 1
                if current_step % self.report_interval == 0:
                    plog.info(''.join([f"Step: {self.global_step}, ", *[f"{k}: {v.value:.4f}, " for k, v in metrics.items()]]))
        
        self.model.save(osp.join(self.trial_save_path, 'model'), self.global_step)

class ABTestTrainer:

    def __init__(
        self, 
        ab_model: AB[Model],
        save_path: str,
        report_interval: int = 100,
        smooth_factor: float = 0.99,
        monitor: AB[Monitor] = AB(Monitor()),
    ):
        self.ab_model = ab_model
        self.save_path = save_path
        self.report_interval = report_interval
        self.metrics_tracker = MetricsTracker(window_size=report_interval)
        self.smooth_factor = smooth_factor
        self.monitor = monitor
        
        os.makedirs(self.trial_save_path)
        self.plog = ProgressLogger(osp.join(self.trial_save_path, "progress.log"))


    @cached_property
    def trial_save_path(self) -> str:
        # in the save_path, there will be subfolders for each run(trial) with autoincrementing index.
        if osp.exists(self.save_path):
            last_index = max([int(trial_folder[6:]) for trial_folder in os.listdir(self.save_path) if trial_folder.startswith("trial_") and trial_folder[6:].isdigit()])
        else:
            last_index = 0
        return osp.join(self.save_path, f"trial_{last_index+1:02d}")
        

    def train_step(self, batch, global_step) -> Optional[Compare]:
        metrics_A = self.ab_model.A.train_on_batch(batch)
        metrics_B = self.ab_model.B.train_on_batch(batch)
        self.metrics_tracker.update(metrics_A, metrics_B)
        self.ab_model.A.state.history[global_step] = deepcopy(metrics_A)
        self.ab_model.B.state.history[global_step] = deepcopy(metrics_B)
        self.monitor.A.log(metrics_A)
        self.monitor.B.log(metrics_B)
        if global_step % self.report_interval == 0:
            return self.metrics_tracker.compare()

    @property
    def global_step(self):
        assert self.ab_model.A.state.global_step == self.ab_model.B.state.global_step
        return self.ab_model.A.state.global_step
    
    @global_step.setter
    def global_step(self, value):
        self.ab_model.A.state.global_step = value
        self.ab_model.B.state.global_step = value

    def train(self, data: StepBasedDataLoader, match_steps, compete_steps, p_thres=0.05):
        plog = self.plog
    
        plog.info(f"Model A: {self.ab_model.A.summary}")
        plog.info(f"Model B: {self.ab_model.B.summary}")
        # Match Stage
        # The test group (B) is assumed to perform worse than the control group (A) initially.
        # We freeze the shared baseline model and train the test group (B) to perform better or equal to the control group (A).
        if match_steps > 0:
            smoothed_report = Compare(p_value=0.0, avg_A=0.6, avg_B=0.4)
            self.ab_model.A.freeze_baseline()
            self.ab_model.B.freeze_baseline()
            with plog.progress_bar(range(match_steps), desc="Match stage") as pbar:
                for current_match_step in pbar:
                    batch = data.get_batch_at_step(self.global_step)
                    report = self.train_step(batch, self.global_step)
                    if report is not None:
                        smoothed_report = report.smoother(
                            smoothed_report, self.smooth_factor)
                        plog.info(f"Mstep: {self.global_step}, ", *smoothed_report.format())
                        if (report.B_better(p_thres) or report.no_difference(p_thres)) and \
                                (smoothed_report.B_better(p_thres) or smoothed_report.no_difference(p_thres)):
                            plog.success(
                                f"Match phase converged in {current_match_step} steps.")
                            break
                    self.global_step += 1
                else:
                    plog.warning(
                        f"Match phase did not converged in {match_steps} steps.")
            self.ab_model.A.unfreeze_baseline()
            self.ab_model.B.unfreeze_baseline()

        # Compete Stage
        # We unfreeze the shared baseline model and train both the control group (A) and test group (B) together
        # to see which will significantly perform better.
        # We assume the initial performance is the same for both groups.
        smoothed_report = Compare(p_value=1.0, avg_A=0.5, avg_B=0.5)
        with plog.progress_bar(range(compete_steps), desc="Compete stage") as pbar:
            for current_compete_step in pbar:
                batch = data.get_batch_at_step(self.global_step)
                report = self.train_step(batch, self.global_step)
                if report is not None:
                    smoothed_report = report.smoother(
                        smoothed_report, self.smooth_factor)
                    plog.info(f"Cstep: {self.global_step}, ", *smoothed_report.format())
                    if (report.B_better(p_thres) and smoothed_report.B_better(p_thres)):
                        plog.success(
                            f"Congratulations! test group (B) is significantly better than the control group (A) after {current_compete_step} steps of training.")
                        break
                    elif (report.A_better(p_thres) and smoothed_report.A_better(p_thres)):
                        plog.info(
                            f"test group (B) performs significantly poorly compared to the control group (A) after {current_compete_step} steps of training.")
                        break
                    else:
                        pass
                self.global_step += 1
            else:
                plog.info(
                    f"There is no significant difference between A and B after {compete_steps} steps of training.")

        # Save the models
        self.ab_model.A.save(osp.join(self.trial_save_path, "A"), self.global_step)
        self.ab_model.B.save(osp.join(self.trial_save_path, "B"), self.global_step)
