import pytorch_lightning as pl
import torch
import torch.nn as nn
import os
import numpy as np
import hydra
from model import load_ssl_model, PhonemeEncoder, DomainEmbedding, LDConditioner, Projection


class BaselineLightningModule(pl.LightningModule):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.construct_model()
        self.save_hyperparameters()
    
    def construct_model(self):
        self.feature_extractors = nn.ModuleList([
            load_ssl_model(cp_path='wav2vec_small.pt'),
            DomainEmbedding(3,128),
        ])
        output_dim = sum([ feature_extractor.get_output_dim() for feature_extractor in self.feature_extractors])
        output_layers = [
            LDConditioner(judge_dim=128,num_judges=3000,input_dim=output_dim)
        ]
        output_dim = output_layers[-1].get_output_dim()
        output_layers.append(
            Projection(hidden_dim=2048,activation=torch.nn.ReLU(),range_clipping=False,input_dim=output_dim)

        )

        self.output_layers = nn.ModuleList(output_layers)

    def forward(self, inputs):
        outputs = {}
        for feature_extractor in self.feature_extractors:
            outputs.update(feature_extractor(inputs))
        x = outputs
        for output_layer in self.output_layers:
            x = output_layer(x,inputs)
        return x