import torch
import torch.nn as nn
from .model import IntentClassifier
from .spacy_loader import nlp
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
from torch.utils.data import TensorDataset, DataLoader
from typing import List, Dict

class IntentModel:
    def __init__(self, data=None, hidden_dims=[128, 64], dropout=0.3):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        """
        data: List of dicts, each dict = {"text": ..., "label": ...}
        If data is None, the model can still be loaded from a saved checkpoint.
        """
        self.hidden_dims = hidden_dims
        self.dropout = dropout

        if data is not None:
            self.data = shuffle(data, random_state=42)
            self.labels = sorted(list({sample['label'] for sample in data}))
            self.label2idx = {label: idx for idx, label in enumerate(self.labels)}
            self.idx2label = {idx: label for label, idx in self.label2idx.items()}
            self.num_classes = len(self.labels)
            self.model = IntentClassifier(num_classes=self.num_classes,
                                         hidden_dims=self.hidden_dims,
                                         dropout=self.dropout).to(self.device)
            self.loss_fn = nn.CrossEntropyLoss()
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=5e-4, weight_decay=1e-4)
            self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer,
                mode='min',
                factor=0.5,
                patience=5,
                min_lr=1e-4
            )

            self._prepare_data()
        else:
            # Placeholder attributes for loading later
            self.data = None
            self.labels = None
            self.label2idx = None
            self.idx2label = None
            self.num_classes = None
            self.model = None
            self.loss_fn = None
            self.optimizer = None
            self.X_train = None
            self.X_val = None
            self.y_train = None
            self.y_val = None
            
    def _prepare_data(self):
        X, y = [], []
        for sample in self.data:
            vec = nlp(sample['text']).vector
            X.append(torch.tensor(vec, dtype=torch.float32))
            y.append(torch.tensor(self.label2idx[sample['label']], dtype=torch.long))
        X = torch.stack(X)
        y = torch.tensor(y)
        self.X_train, self.X_val, self.y_train, self.y_val = train_test_split(
            X.to(self.device), y.to(self.device), test_size=0.2, random_state=42
        )

        train_dataset = TensorDataset(self.X_train, self.y_train)
        self.train_loader = DataLoader(
            train_dataset,
            batch_size=32,
            shuffle=True
        )

    def train(self, epochs=100, verbose=True):
        best_val_loss = float('inf')
        patience_counter = 0
        patience_limit = 8   # stop if val loss doesn't improve for 8 epochs

        for epoch in range(epochs):
            # ---------- TRAIN ----------
            self.model.train()
            total_loss, correct, total = 0.0, 0, 0

            for xb, yb in self.train_loader:
                preds = self.model(xb)
                loss = self.loss_fn(preds, yb)

                self.optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.optimizer.step()

                total_loss += loss.item() * yb.size(0)
                correct += (preds.argmax(1) == yb).sum().item()
                total += yb.size(0)

            train_loss = total_loss / total
            train_acc = correct / total

            # ---------- VALIDATION ----------
            self.model.eval()
            with torch.inference_mode():
                val_preds = self.model(self.X_val)
                val_loss = self.loss_fn(val_preds, self.y_val)
                val_acc = (val_preds.argmax(1) == self.y_val).float().mean().item()

            # ---------- LR SCHEDULER ----------
            self.scheduler.step(val_loss)

            # ---------- EARLY STOPPING ----------
            if val_loss < best_val_loss - 1e-4:  # tiny tolerance
                best_val_loss = val_loss
                patience_counter = 0
            else:
                patience_counter += 1

            # ---------- LOGGING ----------
            if verbose:
                print(
                    f"Epoch {epoch+1}/{epochs} | "
                    f"LR {self.optimizer.param_groups[0]['lr']:.6f} | "
                    f"Train Loss {train_loss:.4f} | Train Acc {train_acc:.2%} | "
                    f"Val Loss {val_loss.item():.4f} | Val Acc {val_acc:.2%}"
                )

            # ---------- STOP ----------
            if patience_counter >= patience_limit:
                print(f"Early stopping triggered at epoch {epoch+1}")
                break

    def finetune(self, new_data: List[Dict], epochs=50, save_path=None):
        """
        Users can finetune with new intents or more data.
        """
        # Initialize data if model was loaded without original data
        if self.data is None:
            self.data = []

        self.data.extend(new_data)
        self.labels = sorted(list({sample['label'] for sample in self.data}))
        self.label2idx = {label: idx for idx, label in enumerate(self.labels)}
        self.idx2label = {idx: label for label, idx in self.label2idx.items()}
        self.num_classes = len(self.labels)

        # Initialize model if missing (loaded empty)
        if self.model is None:
            self.model = IntentClassifier(num_classes=self.num_classes,
                                        hidden_dims=self.hidden_dims,
                                        dropout=self.dropout).to(self.device)

        # Rebuild model if number of classes changed
        elif self.num_classes != self.model.fc[-1].out_features:
            print("Updating model for new number of classes...")
            old_state = self.model.state_dict()
            self.model = IntentClassifier(num_classes=self.num_classes,
                                        hidden_dims=self.hidden_dims,
                                        dropout=self.dropout).to(self.device)
            self.model.load_state_dict({k: v for k, v in old_state.items() if k in self.model.state_dict()})

        # Initialize optimizer and loss function if missing
        if self.optimizer is None:
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=5e-4)

        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            mode='min',
            factor=0.5,
            patience=5,
        )
        if self.loss_fn is None:
            self.loss_fn = nn.CrossEntropyLoss()

        # Prepare data for training
        self._prepare_data()

        # Train
        self.train(epochs=epochs)

        # Save finetuned model if path provided
        if save_path:
            self.save_model(save_path)

    def predict(self, text, threshold=0.6):
        self.model.eval()
        with torch.no_grad(): # Speeds up inference
            vec = torch.tensor(nlp(text).vector, dtype=torch.float32).unsqueeze(0).to(self.device)
            logits = self.model(vec)
            probs = torch.softmax(logits, dim=1)[0]

            max_prob, pred_idx = torch.max(probs, dim=0)

            if max_prob.item() < threshold:
                return "unknown_intent", float(max_prob.item())

            return self.idx2label[pred_idx.item()], float(max_prob.item())

    def load_model(self, path="intent_model.pth"):
        """
        Load model weights + metadata
        """
        checkpoint = torch.load(path, map_location=self.device)
        self.label2idx = checkpoint['label2idx']
        self.idx2label = checkpoint['idx2label']
        self.num_classes = checkpoint['num_classes']
        self.hidden_dims = checkpoint['hidden_dims']
        self.dropout = checkpoint['dropout']
        self.model = IntentClassifier(num_classes=self.num_classes,
                                     hidden_dims=self.hidden_dims,
                                     dropout=self.dropout).to(self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=5e-4, weight_decay=1e-4)
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            mode='min',
            factor=0.5,
            patience=5,
            min_lr=1e-4
        )

        print(f"Model loaded from {path}")

    def save_model(self, path="intent_model.pth"):
        """
        Save the model weights + metadata
        """
        torch.save({
            'model_state_dict': self.model.state_dict(),
            'label2idx': self.label2idx,
            'idx2label': self.idx2label,
            'num_classes': self.num_classes,
            'hidden_dims': self.hidden_dims,
            'dropout': self.dropout
        }, path)
        print(f"Model saved to {path}")

    def load_model(self, path="intent_model.pth"):
        """
        Load model weights + metadata
        """
        checkpoint = torch.load(path, map_location=self.device)
        self.label2idx = checkpoint['label2idx']
        self.idx2label = checkpoint['idx2label']
        self.num_classes = checkpoint['num_classes']
        self.hidden_dims = checkpoint['hidden_dims']
        self.dropout = checkpoint['dropout']
        self.model = IntentClassifier(num_classes=self.num_classes,
                                     hidden_dims=self.hidden_dims,
                                     dropout=self.dropout).to(self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=5e-4, weight_decay=1e-4)
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            mode='min',
            factor=0.5,
            patience=5,
            min_lr=1e-4
        )

        print(f"Model loaded from {path}")