import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import os
from datetime import datetime
import json




def get_model_device(model):
    return next(iter(model.parameters())).device


class RGLRU(nn.Module):
    def __init__(self, hidden_size: int, c: float = 8.0):
        super().__init__()
        self.hidden_size = hidden_size
        self.c = c

        self.input_gate = nn.Linear(hidden_size, hidden_size, bias=False)
        self.recurrence_gate = nn.Linear(hidden_size, hidden_size, bias=False)

        self._base_param = nn.Parameter(torch.empty(hidden_size))
        nn.init.normal_(self._base_param, mean=0.0, std=1.0)  # ok to be any real

    def forward(self, x_t: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
        batch_size, hidden_size = x_t.shape
        assert hidden_size == self.hidden_size
        assert state.shape[0] == batch_size

        i_t = torch.sigmoid(self.input_gate(x_t))
        r_t = torch.sigmoid(self.recurrence_gate(x_t))  # in (0,1)

        eps = 1e-4
        base = torch.sigmoid(self._base_param).unsqueeze(0)  # shape (1, hidden)
        base = base.clamp(min=eps, max=1.0 - eps)

        # exponent = c * r_t (positive)
        a_t = base ** (
            self.c * r_t
        )  # shape (batch, hidden), safe because base in (0,1)

        # ensure numerical stability for sqrt
        one_minus_sq = 1.0 - a_t * a_t
        one_minus_sq = torch.clamp(one_minus_sq, min=0.0)
        multiplier = torch.sqrt(one_minus_sq)

        new_state = (state * a_t) + (multiplier * (i_t * x_t))

        return new_state

    def init_state(self, batch_size: int, device: torch.device | None = None):
        if device is None:
            device = get_model_device(self)
        return torch.zeros(batch_size, self.hidden_size, device=device)


class CausalConv1d(nn.Module):
    def __init__(self, hidden_size, kernel_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.kernel_size = kernel_size
        self.conv = nn.Conv1d(
            hidden_size, hidden_size, kernel_size, groups=hidden_size, bias=True
        )

    def init_state(self, batch_size: int, device: torch.device | None = None):
        if device is None:
            device = get_model_device(self)
        return torch.zeros(
            batch_size, self.hidden_size, self.kernel_size - 1, device=device
        )

    def forward(self, x: torch.Tensor, state: torch.Tensor):
        x_with_state = torch.concat([state, x[:, :, None]], dim=-1)
        out = self.conv(x_with_state)
        new_state = x_with_state[:, :, 1:]
        return out.squeeze(-1), new_state


class Hawk(nn.Module):
    def __init__(self, hidden_size: int, conv_kernel_size: int = 4):
        super().__init__()

        self.conv_kernel_size = conv_kernel_size
        self.hidden_size = hidden_size

        self.gate_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.recurrent_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.conv = CausalConv1d(hidden_size, conv_kernel_size)
        self.rglru = RGLRU(hidden_size)
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False)

    def forward(
        self, x: torch.Tensor, state: tuple[torch.Tensor, torch.Tensor]
    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
        conv_state, rglru_state = state

        batch_size, hidden_size = x.shape
        assert batch_size == conv_state.shape[0] == rglru_state.shape[0]
        assert self.hidden_size == hidden_size == rglru_state.shape[1]

        gate = F.gelu(self.gate_proj(x))
        x = self.recurrent_proj(x)

        x, new_conv_state = self.conv(x, conv_state)
        new_rglru_state = self.rglru(x, rglru_state)

        gated = gate * new_rglru_state
        out = self.out_proj(gated)

        new_state = [new_conv_state, new_rglru_state]
        return out, new_state

    def init_state(
        self, batch_size: int, device: torch.device | None = None
    ) -> list[torch.Tensor]:
        return [
            self.conv.init_state(batch_size, device),
            self.rglru.init_state(batch_size, device),
        ]


class HawkPredictor(nn.Module):
    """Full model with input projection and output head"""

    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 2,
        conv_kernel_size: int = 4,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        # Input projection
        self.input_proj = nn.Linear(input_size, hidden_size)
        self.input_norm = nn.LayerNorm(hidden_size)

        # Hawk layers
        self.hawk_layers = nn.ModuleList(
            [Hawk(hidden_size, conv_kernel_size) for _ in range(num_layers)]
        )

        # Layer norms
        self.layer_norms = nn.ModuleList(
            [nn.LayerNorm(hidden_size) for _ in range(num_layers)]
        )

        # Dropout
        self.dropout = nn.Dropout(dropout)

        # Output head
        self.output_head = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, 1),
        )

    def forward(self, x: torch.Tensor, states=None):
        """
        Args:
            x: (batch_size, seq_len, input_size)
            states: list of states for each layer
        Returns:
            predictions: (batch_size, seq_len, 1)
            final_states: list of final states
        """
        batch_size, seq_len, _ = x.shape
        device = x.device

        # Initialize states if needed
        if states is None:
            states = [
                layer.init_state(batch_size, device) for layer in self.hawk_layers
            ]

        # Input projection
        x = self.input_proj(x)  # (batch, seq, hidden)
        x = self.input_norm(x)

        outputs = []
        final_states = []

        # Process sequence
        for t in range(seq_len):
            x_t = x[:, t, :]  # (batch, hidden)

            # Pass through Hawk layers
            new_states = []
            for i, (hawk_layer, layer_norm) in enumerate(
                zip(self.hawk_layers, self.layer_norms)
            ):
                residual = x_t
                x_t, state = hawk_layer(x_t, states[i])
                x_t = layer_norm(x_t + residual)
                x_t = self.dropout(x_t)
                new_states.append(state)

            states = new_states
            outputs.append(x_t)

        # Stack outputs
        outputs = torch.stack(outputs, dim=1)  # (batch, seq, hidden)

        # Generate predictions
        predictions = self.output_head(outputs)  # (batch, seq, 1)

        return predictions, states




class TimeSeriesDataset(Dataset):
    def __init__(self, features, targets, seq_length=20):
        self.features = features
        self.targets = targets
        self.seq_length = seq_length

    def __len__(self):
        return len(self.features) - self.seq_length

    def __getitem__(self, idx):
        x = self.features[idx : idx + self.seq_length]
        y = self.targets[idx : idx + self.seq_length]
        return torch.FloatTensor(x), torch.FloatTensor(y).squeeze(-1)




class MetricsLogger:
    def __init__(self, save_dir):
        self.save_dir = save_dir
        self.metrics = {
            "train_loss": [],
            "val_loss": [],
            "train_mse": [],
            "val_mse": [],
            "train_mae": [],
            "val_mae": [],
            "learning_rates": [],
        }

    def update(self, epoch_metrics):
        for key, value in epoch_metrics.items():
            if key in self.metrics:
                self.metrics[key].append(value)

    def save(self):
        with open(os.path.join(self.save_dir, "metrics.json"), "w") as f:
            json.dump(self.metrics, f, indent=4)

    def plot_metrics(self):
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle("Training Metrics", fontsize=16)

        # Loss
        ax = axes[0, 0]
        ax.plot(self.metrics["train_loss"], label="Train Loss", marker="o")
        ax.plot(self.metrics["val_loss"], label="Val Loss", marker="s")
        ax.set_xlabel("Epoch")
        ax.set_ylabel("Loss")
        ax.set_title("Training and Validation Loss")
        ax.legend()
        ax.grid(True)

        # MSE
        ax = axes[0, 1]
        ax.plot(self.metrics["train_mse"], label="Train MSE", marker="o")
        ax.plot(self.metrics["val_mse"], label="Val MSE", marker="s")
        ax.set_xlabel("Epoch")
        ax.set_ylabel("MSE")
        ax.set_title("Mean Squared Error")
        ax.legend()
        ax.grid(True)

        # MAE
        ax = axes[1, 0]
        ax.plot(self.metrics["train_mae"], label="Train MAE", marker="o")
        ax.plot(self.metrics["val_mae"], label="Val MAE", marker="s")
        ax.set_xlabel("Epoch")
        ax.set_ylabel("MAE")
        ax.set_title("Mean Absolute Error")
        ax.legend()
        ax.grid(True)

        # Learning Rate
        ax = axes[1, 1]
        ax.plot(self.metrics["learning_rates"], marker="o", color="purple")
        ax.set_xlabel("Epoch")
        ax.set_ylabel("Learning Rate")
        ax.set_title("Learning Rate Schedule")
        ax.grid(True)
        ax.set_yscale("log")

        plt.tight_layout()
        plt.savefig(os.path.join(self.save_dir, "training_metrics.png"), dpi=300)
        plt.close()


def calculate_metrics(predictions, targets):
    """Calculate MSE and MAE"""
    mse = F.mse_loss(predictions, targets).item()
    mae = F.l1_loss(predictions, targets).item()
    return mse, mae


def save_checkpoint(
    model, optimizer, scheduler, epoch, metrics, save_dir, is_best=False
):
    checkpoint = {
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
        "metrics": metrics,
    }

    # Save regular checkpoint
    checkpoint_path = os.path.join(save_dir, f"checkpoint_epoch_{epoch}.pt")
    torch.save(checkpoint, checkpoint_path)

    # Save best model
    if is_best:
        best_path = os.path.join(save_dir, "best_model.pt")
        torch.save(checkpoint, best_path)
        print(f"✓ Saved best model at epoch {epoch}")

    # Keep only last 5 checkpoints
    checkpoints = sorted(
        [f for f in os.listdir(save_dir) if f.startswith("checkpoint_epoch_")]
    )
    if len(checkpoints) > 5:
        for old_ckpt in checkpoints[:-5]:
            os.remove(os.path.join(save_dir, old_ckpt))




def train_epoch(model, train_loader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    all_predictions = []
    all_targets = []

    for batch_idx, (x, y) in enumerate(train_loader):
        x, y = x.to(device), y.to(device)

        optimizer.zero_grad()

        # Forward pass
        predictions, _ = model(x)
        predictions = predictions.squeeze(-1)

        # Calculate loss
        loss = criterion(predictions, y)

        # Backward pass
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        total_loss += loss.item()
        all_predictions.append(predictions.detach())
        all_targets.append(y.detach())

    avg_loss = total_loss / len(train_loader)
    all_predictions = torch.cat(all_predictions, dim=0)
    all_targets = torch.cat(all_targets, dim=0)
    mse, mae = calculate_metrics(all_predictions, all_targets)

    return avg_loss, mse, mae


def validate(model, val_loader, criterion, device):
    model.eval()
    total_loss = 0
    all_predictions = []
    all_targets = []

    with torch.no_grad():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)

            predictions, _ = model(x)
            predictions = predictions.squeeze(-1)

            loss = criterion(predictions, y)

            total_loss += loss.item()
            all_predictions.append(predictions)
            all_targets.append(y)

    avg_loss = total_loss / len(val_loader)
    all_predictions = torch.cat(all_predictions, dim=0)
    all_targets = torch.cat(all_targets, dim=0)
    mse, mae = calculate_metrics(all_predictions, all_targets)

    return avg_loss, mse, mae


def train_model(model, train_loader, val_loader, config):
    """Main training loop"""
    device = config["device"]
    model = model.to(device)

    # Setup
    criterion = nn.MSELoss()
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config["learning_rate"],
        weight_decay=config["weight_decay"],
    )
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=5, verbose=True
    )

    # Create save directory
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    save_dir = os.path.join(config["save_dir"], f"run_{timestamp}")
    os.makedirs(save_dir, exist_ok=True)

    # Save config
    with open(os.path.join(save_dir, "config.json"), "w") as f:
        json.dump(config, f, indent=4)

    # Initialize logger
    logger = MetricsLogger(save_dir)
    best_val_loss = float("inf")

    print(f"{'='*60}")
    print(f"Training started at {timestamp}")
    print(f"Model: {config['model_name']}")
    print(f"Device: {device}")
    print(f"Save directory: {save_dir}")
    print(f"{'='*60}\n")

    # Training loop
    for epoch in range(1, config["num_epochs"] + 1):
        # Train
        train_loss, train_mse, train_mae = train_epoch(
            model, train_loader, optimizer, criterion, device
        )

        # Validate
        val_loss, val_mse, val_mae = validate(model, val_loader, criterion, device)

        # Update scheduler
        scheduler.step(val_loss)
        current_lr = optimizer.param_groups[0]["lr"]

        # Log metrics
        epoch_metrics = {
            "train_loss": train_loss,
            "val_loss": val_loss,
            "train_mse": train_mse,
            "val_mse": val_mse,
            "train_mae": train_mae,
            "val_mae": val_mae,
            "learning_rates": current_lr,
        }
        logger.update(epoch_metrics)

        # Print progress
        print(f"Epoch {epoch}/{config['num_epochs']}")
        print(
            f"  Train - Loss: {train_loss:.6f}, MSE: {train_mse:.6f}, MAE: {train_mae:.6f}"
        )
        print(f"  Val   - Loss: {val_loss:.6f}, MSE: {val_mse:.6f}, MAE: {val_mae:.6f}")
        print(f"  LR: {current_lr:.2e}")

        # Save checkpoint
        is_best = val_loss < best_val_loss
        if is_best:
            best_val_loss = val_loss

        if epoch % config["save_every"] == 0 or is_best:
            save_checkpoint(
                model, optimizer, scheduler, epoch, epoch_metrics, save_dir, is_best
            )

        # Plot metrics every 10 epochs
        if epoch % 10 == 0:
            logger.plot_metrics()

        print()

    # Final save
    logger.save()
    logger.plot_metrics()

    print(f"{'='*60}")
    print(f"Training completed!")
    print(f"Best validation loss: {best_val_loss:.6f}")
    print(f"Results saved to: {save_dir}")
    print(f"{'='*60}")

    return model, logger



if __name__ == "__main__":
    from data_prep.data_clean import clean_indicator
    from data_prep.data_load import prepare_data
    torch.autograd.set_detect_anomaly(True)
    # Configuration
    config = {
      'model_name': 'HawkPredictor',
      'seq_length': 20,
      'hidden_size': 128,
      'num_layers': 3,
      'conv_kernel_size': 4,
      'dropout': 0.2,
      'batch_size': 64,
      'num_epochs': 100,
      'learning_rate': 0.001,
      'weight_decay': 1e-5,
      'train_split': 0.8,
      'save_every': 5,
      'save_dir': './checkpoints',
      'device': 'cuda' if torch.cuda.is_available() else 'cpu'
  }

    print("Loading data...")
    test_dir = "/home/aman/code/ml_fr/ml_stocks/data/NIFTY_5_years.csv"

    load_df = prepare_data(test_dir)
    df = clean_indicator(load_df)

    # Prepare features and target
    target_col = "Daily_Return"
    feature_cols = [col for col in df.columns if col != target_col]

    train_size = int(len(df) * config["train_split"])
    train_df = df[:train_size]
    val_df = df[train_size:]

    scaler = StandardScaler()
    train_features = scaler.fit_transform(train_df[feature_cols].values)
    val_features = scaler.transform(val_df[feature_cols].values)

    train_targets = train_df[target_col].values.reshape(-1, 1)
    val_targets = val_df[target_col].values.reshape(-1, 1)

    # Create datasets
    train_dataset = TimeSeriesDataset(
        train_features, train_targets, config["seq_length"]
    )
    val_dataset = TimeSeriesDataset(val_features, val_targets, config["seq_length"])

    train_loader = DataLoader(
        train_dataset, batch_size=config["batch_size"], shuffle=True, num_workers=0
    )
    val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], 
                          shuffle=False, num_workers=0)

    print(f"Training samples: {len(train_dataset)}")
    print(f"Validation samples: {len(val_dataset)}")
    print(f"Input features: {len(feature_cols)}")

    # Initialize model
    model = HawkPredictor(
      input_size=len(feature_cols),
      hidden_size=config['hidden_size'],
      num_layers=config['num_layers'],
      conv_kernel_size=config['conv_kernel_size'],
      dropout=config['dropout']
  )

    print(f"\nModel parameters: {sum(p.numel() for p in model.parameters()):,}")

    # Train model
    trained_model, metrics_logger = train_model(model, train_loader, val_loader, config)

    print("\nTraining complete! Check the checkpoints directory for saved models and metrics.")