import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import get_linear_schedule_with_warmup
from tqdm import tqdm
import logging
import wandb
from pathlib import Path
from typing import Dict, Any
from torch.cuda.amp import autocast, GradScaler
from datetime import datetime

from data2 import data_processing
from alignment_model import ImageTextAlignmentModel
from report_generator import MedicalReportGenerator
from biovil_t.pretrained import get_biovil_t_image_encoder  # Ensure this import path is correct
from rouge_score import rouge_scorer

def train_epoch(image_encoder, alignment_model, report_generator, train_loader,
                contrastive_loss, alignment_optimizer, generator_optimizer,
                alignment_scheduler, generator_scheduler, scaler, device,
                gradient_accumulation_steps, max_grad_norm, epoch):
    alignment_model.train()
    report_generator.train()
    image_encoder.eval()

    # Metrics tracking
    total_train_loss = 0.0
    total_align_loss = 0.0
    total_gen_loss = 0.0
    total_samples = 0

    progress_bar = tqdm(train_loader, desc=f'Training Epoch {epoch}')

    for batch_idx, (images, findings_texts, findings_lists) in enumerate(progress_bar):
        images = images.to(device)
        batch_size = images.size(0)
        total_samples += batch_size

        # Get image embeddings
        with torch.no_grad():
            image_embeddings = image_encoder(images).img_embedding

        # Create prompts using findings_lists (for generation)
        batch_prompts = [
            f"Findings: {', '.join(findings) if findings else 'No Findings'}."
            for findings in findings_lists
        ]

        # Use findings_texts (actual findings) for alignment
        actual_findings = findings_texts

        # Mixed precision training
        with autocast():
            # Alignment phase
            projected_image, projected_text = alignment_model(image_embeddings, actual_findings)

            # Contrastive loss
            labels = torch.ones(batch_size).to(device)
            align_loss = contrastive_loss(projected_image, projected_text, labels)
            align_loss = align_loss / gradient_accumulation_steps

        # Scale and accumulate alignment gradients
        scaler.scale(align_loss).backward()

        # Generation phase

        # Tokenize the prompts
        prompt_encoding = report_generator.tokenizer(
            batch_prompts,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=512
        ).to(device)

        # Tokenize target texts (actual findings)
        target_encoding = report_generator.tokenizer(
            actual_findings,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=512
        ).to(device)

        with autocast():
            gen_loss, _ = report_generator(
                image_embeddings=image_embeddings.detach(),
                prompt_input_ids=prompt_encoding['input_ids'],
                target_ids=target_encoding['input_ids']
            )
            gen_loss = gen_loss / gradient_accumulation_steps

        # Scale and accumulate generator gradients
        scaler.scale(gen_loss).backward()

        # Update metrics
        total_align_loss += align_loss.item() * gradient_accumulation_steps * batch_size
        total_gen_loss += gen_loss.item() * gradient_accumulation_steps * batch_size
        total_train_loss += (align_loss.item() + gen_loss.item()) * gradient_accumulation_steps * batch_size

        # Step optimizers and schedulers
        if (batch_idx + 1) % gradient_accumulation_steps == 0:
            # Unscale gradients
            scaler.unscale_(alignment_optimizer)
            scaler.unscale_(generator_optimizer)

            # Clip gradients
            torch.nn.utils.clip_grad_norm_(
                alignment_model.parameters(), max_grad_norm
            )
            torch.nn.utils.clip_grad_norm_(
                report_generator.parameters(), max_grad_norm
            )

            # Step optimizers
            scaler.step(alignment_optimizer)
            scaler.step(generator_optimizer)
            scaler.update()

            # Zero gradients
            alignment_optimizer.zero_grad()
            generator_optimizer.zero_grad()

            # Step schedulers
            alignment_scheduler.step()
            generator_scheduler.step()

        # Update progress bar
        progress_bar.set_postfix({
            'align_loss': f"{align_loss.item():.4f}",
            'gen_loss': f"{gen_loss.item():.4f}"
        })

    epoch_align_loss = total_align_loss / total_samples
    epoch_gen_loss = total_gen_loss / total_samples
    epoch_train_loss = total_train_loss / total_samples

    return {
        'train_loss': epoch_train_loss,
        'train_align_loss': epoch_align_loss,
        'train_gen_loss': epoch_gen_loss,
    }

def validate_epoch(image_encoder, alignment_model, report_generator, val_loader,
                   contrastive_loss, device, epoch):
    alignment_model.eval()
    report_generator.eval()
    image_encoder.eval()

    # Metrics storage
    total_val_loss = 0.0
    total_align_loss = 0.0
    total_gen_loss = 0.0
    total_samples = 0
    all_generated = []
    all_references = []

    with torch.no_grad():
        progress_bar = tqdm(val_loader, desc=f'Validation Epoch {epoch}')

        for batch_idx, (images, findings_texts, findings_lists) in enumerate(progress_bar):
            images = images.to(device)
            batch_size = images.size(0)
            total_samples += batch_size

            # Get image embeddings
            image_embeddings = image_encoder(images).img_embedding

            # Create prompts using findings_lists
            batch_prompts = [
                f"Findings: {', '.join(findings) if findings else 'No Findings'}."
                for findings in findings_lists
            ]

            # Actual findings for alignment and reference
            actual_findings = findings_texts

            # Alignment phase
            projected_image, projected_text = alignment_model(image_embeddings, actual_findings)
            labels = torch.ones(batch_size).to(device)
            align_loss = contrastive_loss(projected_image, projected_text, labels)

            # Generation phase
            prompt_encoding = report_generator.tokenizer(
                batch_prompts,
                padding=True,
                truncation=True,
                return_tensors="pt",
                max_length=512
            ).to(device)

            target_encoding = report_generator.tokenizer(
                actual_findings,
                padding=True,
                truncation=True,
                return_tensors="pt",
                max_length=512
            ).to(device)

            # Compute generation loss
            gen_loss, _ = report_generator(
                image_embeddings=image_embeddings,
                prompt_input_ids=prompt_encoding['input_ids'],
                target_ids=target_encoding['input_ids']
            )

            # Generate text for evaluation
            generated_texts = report_generator(
                image_embeddings=image_embeddings,
                prompt_input_ids=prompt_encoding['input_ids'],
                target_ids=None
            )

            # Store the generated and reference texts for ROUGE calculation
            all_generated.extend(generated_texts)
            all_references.extend(actual_findings)

            # Update totals
            total_align_loss += align_loss.item() * batch_size
            total_gen_loss += gen_loss.item() * batch_size
            total_val_loss += (align_loss.item() + gen_loss.item()) * batch_size

            # Print sample generation
            if batch_idx % 10 == 0:
                print(f"\nSample Generation (Batch {batch_idx}):")
                print(f"Generated: {generated_texts[0]}")
                print(f"Reference: {actual_findings[0]}")
                # Also display the pathologies findings from findings_lists
                print(f"Pathologies/Findings List: {findings_lists[0]}\n")

        # Calculate overall metrics
        epoch_align_loss = total_align_loss / total_samples
        epoch_gen_loss = total_gen_loss / total_samples
        epoch_val_loss = total_val_loss / total_samples

    # Compute ROUGE-L
    scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
    rouge_l_scores = []
    for ref, gen in zip(all_references, all_generated):
        score = scorer.score(ref, gen)['rougeL'].fmeasure
        rouge_l_scores.append(score)
    avg_rouge_l = sum(rouge_l_scores) / len(rouge_l_scores) if rouge_l_scores else 0.0

    # Display validation losses and ROUGE-L
    print(f"\nEpoch {epoch} Validation Metrics:")
    print(f"Validation Loss: {epoch_val_loss:.4f}")
    print(f"Alignment Loss: {epoch_align_loss:.4f}")
    print(f"Generation Loss: {epoch_gen_loss:.4f}")
    print(f"ROUGE-L: {avg_rouge_l:.4f}")

    return {
        'val_loss': epoch_val_loss,
        'val_align_loss': epoch_align_loss,
        'val_gen_loss': epoch_gen_loss,
        'val_rouge_l': avg_rouge_l
    }


def train_model(
        csv_with_image_paths: str,
        csv_with_labels: str,
        num_epochs: int = 30,
        batch_size: int = 8,
        train_split: float = 0.85,
        num_workers: int = 4,
        learning_rate: float = 2e-4,
        warmup_steps: int = 1000,
        gradient_accumulation_steps: int = 4,
        max_grad_norm: float = 1.0,
        use_wandb: bool = True,
        checkpoint_dir: str = "checkpoints",
        seed: int = 42
):
    # Set device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Initialize models
    image_encoder = get_biovil_t_image_encoder()
    alignment_model = ImageTextAlignmentModel(image_embedding_dim=512)
    report_generator = MedicalReportGenerator(image_embedding_dim=512)

    # Move models to device
    image_encoder = image_encoder.to(device)
    alignment_model = alignment_model.to(device)
    report_generator = report_generator.to(device)

    # Initialize wandb
    if use_wandb:
        wandb.init(
            project="medical-report-generation",
            config={
                "learning_rate": learning_rate,
                "epochs": num_epochs,
                "batch_size": batch_size,
                "warmup_steps": warmup_steps,
                "gradient_accumulation_steps": gradient_accumulation_steps,
            }
        )
        wandb.watch(models=[alignment_model, report_generator], log="all")

    # Get dataloaders
    train_loader, val_loader = data_processing.get_dataloaders(
        csv_with_image_paths=csv_with_image_paths,
        csv_with_labels=csv_with_labels,
        batch_size=batch_size,
        train_split=train_split,
        num_workers=num_workers,
        seed=seed,
    )

    # Initialize optimizers
    alignment_optimizer = AdamW(
        alignment_model.parameters(),
        lr=learning_rate,
        weight_decay=0.01
    )
    generator_optimizer = AdamW([
        {'params': report_generator.model.parameters(), 'lr': learning_rate},
        {'params': report_generator.image_projection.parameters(), 'lr': learning_rate * 10}
    ])

    # Initialize schedulers
    num_training_steps = len(train_loader) * num_epochs // gradient_accumulation_steps
    alignment_scheduler = get_linear_schedule_with_warmup(
        alignment_optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=num_training_steps
    )
    generator_scheduler = get_linear_schedule_with_warmup(
        generator_optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=num_training_steps
    )

    # Initialize loss function and scaler
    contrastive_loss = nn.CosineEmbeddingLoss()
    scaler = GradScaler()

    # Create checkpoint directory
    checkpoint_dir = Path(checkpoint_dir)
    checkpoint_dir.mkdir(parents=True, exist_ok=True)

    for epoch in range(num_epochs):
        print(f"\nEpoch {epoch + 1}/{num_epochs}")

        # Training phase
        train_metrics = train_epoch(
            image_encoder=image_encoder,
            alignment_model=alignment_model,
            report_generator=report_generator,
            train_loader=train_loader,
            contrastive_loss=contrastive_loss,
            alignment_optimizer=alignment_optimizer,
            generator_optimizer=generator_optimizer,
            alignment_scheduler=alignment_scheduler,
            generator_scheduler=generator_scheduler,
            scaler=scaler,
            device=device,
            gradient_accumulation_steps=gradient_accumulation_steps,
            max_grad_norm=max_grad_norm,
            epoch=epoch + 1
        )

        # Validation phase
        val_metrics = validate_epoch(
            image_encoder=image_encoder,
            alignment_model=alignment_model,
            report_generator=report_generator,
            val_loader=val_loader,
            contrastive_loss=contrastive_loss,
            device=device,
            epoch=epoch + 1
        )

        # Display training and validation losses
        print(f"\nEpoch {epoch + 1} Training Loss: {train_metrics['train_loss']:.4f}")
        print(f"Epoch {epoch + 1} Validation Loss: {val_metrics['val_loss']:.4f}")
        print(f"Alignment Loss - Train: {train_metrics['train_align_loss']:.4f}, Val: {val_metrics['val_align_loss']:.4f}")
        print(f"Generation Loss - Train: {train_metrics['train_gen_loss']:.4f}, Val: {val_metrics['val_gen_loss']:.4f}")
        print(f"ROUGE-L (Val): {val_metrics['val_rouge_l']:.4f}")

        # Log metrics to wandb
        if use_wandb:
            wandb.log({**train_metrics, **val_metrics})

        # Save model checkpoint after each epoch
        checkpoint_save_path = checkpoint_dir / f"model_epoch_{epoch+1}.pt"
        torch.save({
            'epoch': epoch + 1,
            'image_encoder_state_dict': image_encoder.state_dict(),
            'alignment_model_state_dict': alignment_model.state_dict(),
            'report_generator_state_dict': report_generator.state_dict(),
            'alignment_optimizer_state_dict': alignment_optimizer.state_dict(),
            'generator_optimizer_state_dict': generator_optimizer.state_dict(),
            'alignment_scheduler_state_dict': alignment_scheduler.state_dict(),
            'generator_scheduler_state_dict': generator_scheduler.state_dict(),
            'scaler_state_dict': scaler.state_dict(),
            'config': {
                'learning_rate': learning_rate,
                'batch_size': batch_size,
                'gradient_accumulation_steps': gradient_accumulation_steps,
                'max_grad_norm': max_grad_norm,
            }
        }, checkpoint_save_path)
        logging.info(f"Saved checkpoint: {checkpoint_save_path}")

    if use_wandb:
        wandb.finish()


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )

    # Path to your CSV files
    csv_with_image_paths = "/home/ubuntu/NLP/NLP_Project/Temp_3_NLP/Data/final.csv"
    csv_with_labels = "/home/ubuntu/NLP/NLP_Project/Temp_3_NLP/Data/labeled_reports_with_images.csv"

    # Training configuration
    config = {
        'num_epochs': 30,
        'batch_size': 8,
        'learning_rate': 1e-4,
        'warmup_steps': 1000,
        'gradient_accumulation_steps': 4,
        'use_wandb': True,
        'checkpoint_dir': 'checkpoints',
        'seed': 42
    }

    # Start training
    train_model(
        csv_with_image_paths=csv_with_image_paths,
        csv_with_labels=csv_with_labels,
        **config
    )