import torch
import torch.nn as nn
import torch.optim as optim
from typing import Tuple, List, Optional, Callable
import logging
import time
from NeuroFlex.utils.utils import normalize_data
from NeuroFlex.utils.descriptive_statistics import preprocess_data

logging.basicConfig(level=logging.INFO)

class LSTMModule(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, num_layers: int = 1, dropout: float = 0.0, learning_rate: float = 0.001):
        super(LSTMModule, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout_rate = dropout
        self.learning_rate = learning_rate
        self.performance_threshold = 0.8
        self.update_interval = 86400  # 24 hours in seconds
        self.gradient_norm_threshold = 10
        self.performance_history_size = 100

        self.is_trained = False
        self.performance = 0.0
        self.last_update = 0
        self.gradient_norm = 0
        self.performance_history = []

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
        self.to(self.device)

    def forward(self, inputs, initial_state=None):
        if initial_state is None:
            initial_state = self.initialize_state(inputs.size(0))

        outputs, (final_h, final_c) = self.lstm(inputs, initial_state)
        return outputs, (final_h, final_c)

    def initialize_state(self, batch_size):
        return (torch.zeros(self.num_layers, batch_size, self.hidden_size, device=self.device),
                torch.zeros(self.num_layers, batch_size, self.hidden_size, device=self.device))

    def diagnose(self):
        issues = []
        if not self.is_trained:
            issues.append("Model is not trained")
        if self.performance < self.performance_threshold:
            issues.append("Model performance is below threshold")
        if time.time() - self.last_update > self.update_interval:
            issues.append("Model hasn't been updated in 24 hours")
        if self.gradient_norm > self.gradient_norm_threshold:
            issues.append("Gradient explosion detected")
        if len(self.performance_history) > 5 and all(p < 0.01 for p in self.performance_history[-5:]):
            issues.append("Model is stuck in local minimum")
        return issues

    def heal(self, issues):
        for issue in issues:
            if issue == "Model is not trained":
                logging.info("Model needs training")
            elif issue == "Model performance is below threshold":
                self.improve_model()
            elif issue == "Model hasn't been updated in 24 hours":
                self.update_model()
            elif issue == "Gradient explosion detected":
                self.handle_gradient_explosion()
            elif issue == "Model is stuck in local minimum":
                self.escape_local_minimum()

    def improve_model(self):
        logging.info("Improving model performance...")
        self.performance = min(self.performance * 1.2 + 0.01, 1.0)
        self.update_performance()

    def update_model(self):
        logging.info("Updating model...")
        self.last_update = time.time()
        self.update_performance()

    def handle_gradient_explosion(self):
        logging.info("Handling gradient explosion...")
        self.learning_rate *= 0.5

    def escape_local_minimum(self):
        logging.info("Attempting to escape local minimum...")
        self.learning_rate = min(self.learning_rate * 2.5, 0.1)
        logging.info(f"New learning rate: {self.learning_rate}")

    def update_performance(self):
        self.performance_history.append(self.performance)
        if len(self.performance_history) > self.performance_history_size:
            self.performance_history.pop(0)

    def adjust_learning_rate(self):
        if len(self.performance_history) >= 2:
            current_performance = self.performance_history[-1]
            previous_performance = self.performance_history[-2]

            if current_performance > previous_performance:
                self.learning_rate *= 1.05
            elif current_performance < previous_performance:
                self.learning_rate *= 0.95
        else:
            self.learning_rate *= 1.01

        self.learning_rate = max(min(self.learning_rate, 0.1), 1e-5)
        return self.learning_rate

def create_lstm_model(input_size: int, hidden_size: int, num_layers: int = 1, dropout: float = 0.0, learning_rate: float = 0.001) -> LSTMModule:
    return LSTMModule(input_size, hidden_size, num_layers, dropout, learning_rate)

def train_lstm_model(model: LSTMModule,
                     x_train: torch.Tensor,
                     y_train: torch.Tensor,
                     epochs: int = 10,
                     batch_size: int = 32,
                     validation_data: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
                     callback: Optional[Callable[[float], None]] = None) -> dict:
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=model.learning_rate)

    history = {'train_loss': [], 'val_loss': []}

    num_samples = x_train.shape[0]
    num_batches = max(1, num_samples // batch_size)

    logging.info(f"LSTM model initial parameters: {sum(p.numel() for p in model.parameters())}")

    for epoch in range(epochs):
        model.train()
        epoch_loss = 0.0
        for i in range(0, num_samples, batch_size):
            batch_x = x_train[i:i+batch_size].to(model.device)
            batch_y = y_train[i:i+batch_size].to(model.device)

            optimizer.zero_grad()
            outputs, _ = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()

        avg_loss = epoch_loss / num_batches
        logging.info(f"LSTM - Epoch [{epoch+1}/{epochs}], Loss: {avg_loss:.4f}")

        if callback:
            callback(avg_loss)

        model.gradient_norm = sum(p.grad.norm().item() for p in model.parameters() if p.grad is not None)
        model.performance = 1.0 - avg_loss  # Simple performance metric
        model.update_performance()
        model.adjust_learning_rate()

        issues = model.diagnose()
        if issues:
            logging.info(f"Diagnosed issues: {issues}")
            model.heal(issues)

        # Compute epoch loss
        model.eval()
        with torch.no_grad():
            train_outputs, _ = model(x_train)
            train_loss = criterion(train_outputs, y_train).item()
            history['train_loss'].append(train_loss)

            if validation_data:
                val_x, val_y = validation_data
                val_outputs, _ = model(val_x)
                val_loss = criterion(val_outputs, val_y).item()
                history['val_loss'].append(val_loss)

    model.is_trained = True
    model.last_update = time.time()
    logging.info(f"LSTM model final parameters: {sum(p.numel() for p in model.parameters())}")

    return history

def lstm_predict(model: LSTMModule, x: torch.Tensor) -> torch.Tensor:
    model.eval()
    with torch.no_grad():
        outputs, _ = model(x.to(model.device))
        return outputs
