import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertForSequenceClassification
import pandas as pd
from tqdm import tqdm
import os
import time  # ⏱️ Import time for time estimation

# ✅ Configuration
EPOCHS = 3  # 🔹 Change number of epochs
BATCH_SIZE = 16  # 🔹 Change batch size if needed

# ✅ Detect GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
if torch.cuda.is_available():
    print(f"✅ Using GPU: {torch.cuda.get_device_name(0)}")
else:
    print("❌ No GPU found. Using CPU instead.")
print(f"Using device: {device}")

# ✅ Yelp Dataset Class (For AI Response Generation)
class YelpResponseDataset(Dataset):
    def __init__(self, file_path, tokenizer, max_length=256):
        self.data = pd.read_csv(file_path)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        text = self.data.iloc[idx]["text"]
        response = self.data.iloc[idx]["response"]

        tokens = self.tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt",
        )

        return tokens["input_ids"].squeeze(0), tokens["attention_mask"].squeeze(0), response

# ✅ Model for Response Prediction
class ResponseModel(nn.Module):
    def __init__(self):
        super(ResponseModel, self).__init__()
        self.bert = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=768)
    
    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        return outputs.logits

# ✅ Training Function with Time Estimation
def train_response():
    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
    train_dataset = YelpResponseDataset("../data/train.csv", tokenizer)
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

    model = ResponseModel().to(device)
    optimizer = optim.AdamW(model.parameters(), lr=3e-5)
    criterion = nn.MSELoss()

    os.makedirs("models", exist_ok=True)
    model_path = "models/response_model.pth"

    # ✅ Load previous model weights if available
    if os.path.exists(model_path):
        print(f"🔄 Loading previous model weights from {model_path}...")
        model.load_state_dict(torch.load(model_path, map_location=device))
    else:
        print("🚀 Starting training from scratch.")

    model.train()

    # ✅ Estimate total time before training starts
    sample_batch = next(iter(train_loader))
    start_time_test = time.time()
    with torch.no_grad():
        _ = model(sample_batch[0].to(device), sample_batch[1].to(device))
    estimated_batch_time = time.time() - start_time_test
    estimated_time_per_epoch = estimated_batch_time * len(train_loader)
    total_estimated_time = estimated_time_per_epoch * EPOCHS
    print(f"⏳ Estimated total training time: {total_estimated_time / 60:.2f} min")

    total_start_time = time.time()  # ⏱️ Start overall training timer

    for epoch in range(EPOCHS):
        epoch_start_time = time.time()  # ⏱️ Start epoch timer
        total_loss = 0

        for input_ids, attention_mask, response in tqdm(train_loader, desc=f"Training Epoch {epoch+1}/{EPOCHS}"):
            input_ids, attention_mask = input_ids.to(device), attention_mask.to(device)

            optimizer.zero_grad()
            response_pred = model(input_ids, attention_mask)

            # Dummy loss calculation (since AI-generated responses are complex)
            loss = criterion(response_pred, torch.zeros_like(response_pred))  
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        epoch_time = time.time() - epoch_start_time  # ⏱️ Epoch duration
        avg_epoch_time = (time.time() - total_start_time) / (epoch + 1)  # ⏱️ Average per epoch
        remaining_time = avg_epoch_time * (EPOCHS - (epoch + 1))  # ⏱️ Estimate remaining time

        print(f"🕒 Epoch {epoch+1} Time: {epoch_time:.2f} sec | Remaining Time: {remaining_time/60:.2f} min")
        print(f"🎯 Epoch {epoch+1}, Response Model Loss: {total_loss:.4f}")

        # ✅ Save model after every epoch
        torch.save(model.state_dict(), model_path)
        print(f"✅ Model saved to '{model_path}' (Epoch {epoch+1})")

    total_training_time = time.time() - total_start_time  # ⏱️ Total training duration
    print(f"🏁 Training complete in {total_training_time/60:.2f} min!")

if __name__ == "__main__":
    train_response()