# model.py

import torch
import torch.nn as nn
from transformers import AutoModelForSeq2SeqLM
from torchvision import models

class ImageToTextProjector(nn.Module):
    def __init__(self, image_embedding_dim, text_embedding_dim):
        super(ImageToTextProjector, self).__init__()
        self.fc = nn.Linear(image_embedding_dim, text_embedding_dim)
        self.activation = nn.ReLU()
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        x = self.fc(x)
        x = self.activation(x)
        x = self.dropout(x)
        return x

class CombinedModel(nn.Module):
    def __init__(self, video_model, report_generator, num_classes, projector):
        super(CombinedModel, self).__init__()
        self.video_model = video_model
        self.report_generator = report_generator
        self.classifier = nn.Linear(512, num_classes)
        self.projector = projector
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, images, labels=None):
        video_embeddings = self.video_model(images)
        video_embeddings = self.dropout(video_embeddings)
        class_outputs = self.classifier(video_embeddings)
        projected_embeddings = self.projector(video_embeddings)
        encoder_inputs = projected_embeddings.unsqueeze(1)

        if labels is not None:
            outputs = self.report_generator(
                inputs_embeds=encoder_inputs,
                labels=labels
            )
            gen_loss = outputs.loss
            generated_report = None
        else:
            generated_report_ids = self.report_generator.generate(
                inputs_embeds=encoder_inputs,
                max_length=512,
                num_beams=4,
                early_stopping=True
            )
            generated_report = report_generator_tokenizer.batch_decode(
                generated_report_ids, skip_special_tokens=True
            )
            gen_loss = None

        return class_outputs, generated_report, gen_loss