from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, cast, Generator, List

import pandas as pd

from pt_strategy.trading_pair import TradingPair

class Prediction:
    tstamp_: pd.Timestamp
    disequilibrium_: float
    scaled_disequilibrium_: float

    def __init__(self, tstamp: pd.Timestamp, disequilibrium: float, scaled_disequilibrium: float):
        self.tstamp_ = tstamp
        self.disequilibrium_ = disequilibrium
        self.scaled_disequilibrium_ = scaled_disequilibrium

    def to_dict(self) -> Dict[str, Any]:
        return {
            "tstamp": self.tstamp_,
            "disequilibrium": self.disequilibrium_,
            "signed_scaled_disequilibrium": self.scaled_disequilibrium_,
            "scaled_disequilibrium": abs(self.scaled_disequilibrium_),
            # "pair": self.pair_,
        }
    def to_df(self) -> pd.DataFrame:
        return pd.DataFrame([self.to_dict()])
    
class PairsTradingModel(ABC):

    @abstractmethod
    def predict(self, pair: TradingPair) -> Prediction:
        ...
    
    @staticmethod
    def create(config: Dict[str, Any]) -> PairsTradingModel:
        import importlib
        
        model_class_name = config.get("model_class", None)
        assert model_class_name is not None
        module_name, class_name = model_class_name.rsplit(".", 1)
        module = importlib.import_module(module_name)
        model_object = getattr(module, class_name)()
        return cast(PairsTradingModel, model_object)