import pickle
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
from typing import Tuple, Iterator, Dict, Any, Optional
from datetime import datetime, timedelta
from rich.console import Console

# Custom imports
from optrade.data.thetadata.contracts import Contract
from optrade.utils.data.error import DataValidationError, MARKET_HOLIDAY, WEEKEND
from optrade.utils.data.pathing import set_contract_dir

SCRIPT_DIR = Path(__file__).resolve().parent

class ContractDataset:
    """
    A dataset containing options contracts generated with consistent parameters.

    Contracts are generated by starting from total_start_date and advancing by
    contract_stride days until reaching the last valid date that allows for
    contracts within the specified time-to-expiration tolerance.
    """

    def __init__(
        self,
        root: str = "AAPL",
        total_start_date: str = "20231107",
        total_end_date: str = "20241114",
        contract_stride: int = 5,
        interval_min: int = 1,
        right: str = "C",
        target_tte: int = 30,
        tte_tolerance: Tuple[int, int] = (25, 35),
        moneyness: str = "OTM",
        target_band: float = 0.05,
        volatility_scaled: bool = True,
        volatility_scalar: float = 1.0,
        hist_vol: Optional[float] = None,
        verbose: bool = False,
        save_dir: Optional[str] = None
    ) -> None:
        """
        Initialize the ContractDataset with the specified parameters.

        Args:
            root: The security root symbol
            total_start_date: Start date for the dataset (YYYYMMDD)
            total_end_date: End date for the dataset (YYYYMMDD)
            contract_stride: Days between consecutive contracts
            interval_min: Data interval in minutes
            right: Option type (C/P)
            target_tte: Target time to expiration in days
            tte_tolerance: Acceptable range for TTE as (min_days, max_days)
            moneyness: Contract moneyness (OTM/ATM/ITM)
            target_band: Target percentage band for strike selection
            volatility_scaled: Whether to scale by volatility
            volatility_scalar: Scaling factor for volatility
            hist_vol: Historical volatility for dynamic strike selection
            verbose: Whether to print verbose output
        """
        self.root = root
        self.total_start_date = total_start_date
        self.total_end_date = total_end_date
        self.contract_stride = contract_stride
        self.interval_min = interval_min
        self.right = right
        self.target_tte = target_tte
        self.tte_tolerance = tte_tolerance
        self.moneyness = moneyness
        self.target_band = target_band
        self.volatility_scaled = volatility_scaled
        self.volatility_scalar = volatility_scalar
        self.hist_vol = hist_vol
        self.verbose = verbose

        self.ctx = Console()
        self.contracts = []
        self.contract_dir = set_contract_dir(
            SCRIPT_DIR=SCRIPT_DIR,
            root=root,
            start_date=total_start_date,
            end_date=total_end_date,
            contract_stride=contract_stride,
            interval_min=interval_min,
            right=right,
            target_tte=target_tte,
            tte_tolerance=tte_tolerance,
            moneyness=moneyness,
            target_band=target_band,
            volatility_scaled=volatility_scaled,
            volatility_scalar=volatility_scalar,
            hist_vol=hist_vol,
            save_dir=save_dir
        )

    def generate_contracts(self) -> "ContractDataset":
        """
        Generate all contracts in the dataset based on configuration parameters.
        """
        # Parse dates
        start_date = datetime.strptime(self.total_start_date, "%Y%m%d")
        end_date = datetime.strptime(self.total_end_date, "%Y%m%d")
        max_tte = max(self.tte_tolerance)

        # Calculate the latest possible start date
        latest_start = end_date - timedelta(days=max_tte)

        # Generate contracts
        current_date = start_date

        while current_date <= latest_start:
            # Format initial date string
            date_str = current_date.strftime("%Y%m%d")
            attempt_date = current_date
            contract = None

            # Find a valid contract for the current date. Some dates may be ineligible due to holidays or weekends.
            while contract is None and attempt_date <= latest_start:
                attempt_date_str = attempt_date.strftime("%Y%m%d")
                try:
                    contract = Contract.find_optimal(
                        root=self.root,
                        start_date=attempt_date_str,
                        interval_min=self.interval_min,
                        right=self.right,
                        target_tte=self.target_tte,
                        tte_tolerance=self.tte_tolerance,
                        moneyness=self.moneyness,
                        target_band=self.target_band,
                        hist_vol=self.hist_vol,
                        volatility_scaled=self.volatility_scaled,
                        volatility_scalar=self.volatility_scalar,
                    )

                    if attempt_date > current_date:
                        self.ctx.log(f"Found valid contract at shifted date: {attempt_date_str}") if self.verbose else None

                except DataValidationError as e:
                    if e.error_code == WEEKEND:
                        self.ctx.log(f"Skipping weekend: {attempt_date_str}") if self.verbose else None
                        attempt_date += timedelta(days=1)
                    elif e.error_code == MARKET_HOLIDAY:
                        self.ctx.log(f"Skipping market holiday: {attempt_date_str}") if self.verbose else None
                        attempt_date += timedelta(days=1)
                    else:
                        self.ctx.log(f"Unkown error: {str(e)}. Skipping date: {attempt_date_str}.") if self.verbose else None

                    # Check if we've run out of valid dates
                    if attempt_date > latest_start:
                        self.ctx.log(f"Unable to find valid contract starting from {date_str}") if self.verbose else None
                        break

                    continue

            # If we found a valid contract, add it and advance by stride
            if contract is not None:
                self.contracts.append(contract)
                self.ctx.log(f"Added contract: {contract}") if self.verbose else None
                current_date = attempt_date + timedelta(days=self.contract_stride)
            else:
                # If no contract was found, advance by one day to try the next period
                current_date += timedelta(days=1)

            self.ctx.log(f"Next start date: {current_date.strftime('%Y%m%d')}") if self.verbose else None

        return self

    def __len__(self) -> int:
        """Get the number of contracts in the dataset."""
        return len(self.contracts)

    def __getitem__(self, idx) -> Contract:
        """Get a contract by index."""
        return self.contracts[idx]

    def __iter__(self) -> Iterator:
        """Iterate through contracts."""
        return iter(self.contracts)

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert the dataset to a dictionary suitable for serialization.
        Separates initialization parameters from computed attributes.
        """
        return {
            "params": {
                "root": self.root,
                "total_start_date": self.total_start_date,
                "total_end_date": self.total_end_date,
                "contract_stride": self.contract_stride,
                "interval_min": self.interval_min,
                "right": self.right,
                "target_tte": self.target_tte,
                "tte_tolerance": self.tte_tolerance,
                "moneyness": self.moneyness,
                "target_band": self.target_band,
                "volatility_scaled": self.volatility_scaled,
                "volatility_scalar": self.volatility_scalar,
                "hist_vol": self.hist_vol
            },
            "contracts": [contract.model_dump() for contract in self.contracts]
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ContractDataset":
        """
        Create a new dataset instance from a dictionary.
        Handles both initialization parameters and computed attributes.
        """
        # Create instance with initialization parameters
        instance = cls(**data["params"])

        # Restore contracts
        instance.contracts = [Contract(**contract_dict) for contract_dict in data["contracts"]]

        return instance

    def save(self, filename: Optional[str] = None) -> None:
        """
        Save the dataset to a pickle file.

        Args:
            filepath: Optional custom filepath. If None, generates default name

        Returns:
            str: Path where the pickle file was saved
        """
        self.contract_dir.mkdir(parents=True, exist_ok=True)
        filepath = self.contract_dir / "contracts.pkl" if filename is None else self.contract_dir / filename

        # Convert to dictionary and save
        data = self.to_dict()
        with open(filepath, 'wb') as f:
            pickle.dump(data, f)

        self.ctx.log(f"Contract dataset saved to \"{filepath}\"") if self.verbose else None

    @classmethod
    def load(cls, contract_dir: Path, filename: Optional[str]=None) -> "ContractDataset":
        """
        Load a dataset from a pickle file.

        Args:
            filepath: Path to the pickle file

        Returns:
            ContractDataset: Reconstructed dataset with all contracts
        """
        filepath = contract_dir / "contracts.pkl" if filename is None else contract_dir / filename

        with open(filepath, 'rb') as f:
            data = pickle.load(f)

        instance = cls.from_dict(data)
        instance.ctx.log(f"Contract dataset loaded from {filepath}") if instance.verbose else None
        return instance

class ForecastingDataset(Dataset):
    """

    Args:
        data (torch.Tensor): Tensor of shape (num_time_steps, num_features)
        seq_len (int): Length of the lookback window
        pred_len (int): Length of the forecast window
        target_channels (list): Channels to forecast. By default target_channels=[0], which corresponds to the
                                option midprice returns. If None, all channels will be returned for the target.
        dtype (str): Desired data type of the tensor.

    __getitem__:
        Returns:
            input (torch.Tensor): Lookback window of shape (num_features, seq_len)
            target (torch.Tensor): Target window of shape (len(target_channels), pred_len)
    """

    def __init__(self, data, seq_len, pred_len, target_channels=[0], dtype="float32"):
        self.dtype = eval("torch." + dtype)
        if not torch.is_tensor(data):
            self.data = torch.from_numpy(data).type(self.dtype)
        else:
            self.data = data.type(self.dtype)
        self.seq_len = seq_len
        self.pred_len = pred_len

        if target_channels != []:
            self.target_channels = target_channels

    def __len__(self):
        return self.data.shape[0] - self.seq_len - self.pred_len

    def __getitem__(self, idx):
        input = self.data[idx:idx+self.seq_len]

        if self.target_channels:
            target = self.data[idx+self.seq_len:idx+self.seq_len+self.pred_len, self.target_channels]
        else:
            target = self.data[idx+self.seq_len:idx+self.seq_len+self.pred_len]

        return input.transpose(0,1), target.transpose(0,1)

class IntradayForecastingDataset(Dataset):
    """
    A standard forecasting dataset class for PyTorch.

    Args:
        data (torch.Tensor): The time series data in a tensor of shape (num_channels, num_time_steps).
        seq_len (int): The length of the input window.
        pred_len (int): The length of the forecast window.
        target_channels (list): The channels to forecast. If None, all channels are forecasted.
        dtype (str): The datatype of the tensor.

    __getitem__ method:

        Args:
            idx (int): The index of the input window in the time series.
        Returns:
            input_data (torch.Tensor): The lookback window of length seq_len for the given index.
            target_data (torch.Tensor): The forecast window for the given index (continuation of input_data
                                        shifted by pred_len)

    """

    def __init__(self, inputs, targets, dtype="float32"):

        # Convert the data to a tensor and set the datatype
        dtype = eval("torch." + dtype)
        if not torch.is_tensor(inputs):
            self.inputs = torch.from_numpy(inputs).type(dtype)
            self.targets = torch.from_numpy(targets).type(dtype)
        else:
            self.inputs = inputs.type(dtype)
            self.targets = targets.type(dtype)

    def __len__(self):
        return self.inputs.shape[0]

    def __getitem__(self, idx):
        return self.input[idx], self.target[idx]

if __name__=="__main__":

    # Define the parameters for the dataset
    root="AMZN"
    total_start_date="20220101"
    total_end_date="20220301"
    hist_vol = 0.1
    contract_stride=1
    interval_min=1
    right="P"
    target_tte=3
    tte_tolerance = (1,10)
    moneyness = "ITM"
    target_band = 0.05
    volatility_scaled = True
    volatility_scalar = 1.0
    save_dir = None

    contracts = ContractDataset(root=root,
    total_start_date=total_start_date,
    total_end_date=total_end_date,
    contract_stride=contract_stride,
    interval_min=interval_min,
    right=right,
    target_tte=target_tte,
    tte_tolerance=tte_tolerance,
    moneyness=moneyness,
    target_band=target_band,
    volatility_scaled=volatility_scaled,
    volatility_scalar=volatility_scalar,
    hist_vol=hist_vol,
    verbose=True,
    save_dir=save_dir)

    # Generate contracts
    contracts.generate_contracts()

    # Save the contracts
    contracts.save()
    print(f"Saving contracts to {contracts.contract_dir}")

    # Load the generated contracts into a fresh dataset
    contract_dir = set_contract_dir(
        SCRIPT_DIR=SCRIPT_DIR,
        root=root,
        start_date=total_start_date,
        end_date=total_end_date,
        contract_stride=contract_stride,
        interval_min=interval_min,
        right=right,
        target_tte=target_tte,
        tte_tolerance=tte_tolerance,
        moneyness=moneyness,
        target_band=target_band,
        volatility_scaled=volatility_scaled,
        volatility_scalar=volatility_scalar,
        hist_vol=hist_vol,
        save_dir=save_dir
    )
    dataset = ContractDataset.load(contract_dir)
    print(f"Loaded contracts from {dataset.contract_dir}")
