# train_duration_predictor_context.py
# ==============================================================================
# Training Script for Duration Predictor with Context
#
# This script trains ONLY the Duration Predictor:
#   - Predicts phoneme durations based on text and context
#   - Uses PredictorDurationSpeakerContext
#
# Validation performs full OOD inference:
#   1. Predict duration (using this model)
#   2. Predict flow latents (using frozen FSQFlowPredictorConvNeXt)
#   3. Decode audio (using frozen Codec)
# ==============================================================================

import os
import os.path as osp
import shutil
import sys
import time
import yaml
import warnings
import logging
import math
import numpy as np
from typing import Optional

import click
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter

# Fix for "Too many open files" error
torch.multiprocessing.set_sharing_strategy('file_system')

from accelerate import Accelerator, DistributedDataParallelKwargs
from accelerate.logging import get_logger

from fsq_flow_dataset import build_fsq_flow_dataloader
from utils import get_data_path_list

# Import our predictors
from Modules.predictor_fsq_flow_convnext import FSQFlowPredictorConvNeXt
from Modules.predictor_duration_speaker_context import PredictorDurationSpeakerContext

# Import codec for validation
from Modules.codec_decoder_mel_speaker import MelCodecVocoderSpeaker

warnings.simplefilter("ignore")
logger = get_logger(__name__, log_level="DEBUG")

class Optimizer:
    """Simple optimizer wrapper with linear decay from 1e-5 to 2e-5 over 800k steps."""
    def __init__(self, parameters, lr=1e-5, betas=(0.9, 0.999), weight_decay=0.01):
        self.optimizer = torch.optim.AdamW(
            parameters, lr=lr, betas=betas, weight_decay=weight_decay
        )
        self.scheduler = None
        
    def set_scheduler(self, total_steps=800000, start_lr=1e-5, final_lr=2e-5):
        from torch.optim.lr_scheduler import LambdaLR
        
        def lr_lambda(step):
            if step >= total_steps:
                return final_lr / start_lr
            # Linear decay from start_lr to final_lr
            return 1.0 + (final_lr / start_lr - 1.0) * (step / total_steps)
        
        self.scheduler = LambdaLR(self.optimizer, lr_lambda)
    
    def step(self):
        self.optimizer.step()
        if self.scheduler is not None:
            self.scheduler.step()
    
    def zero_grad(self):
        self.optimizer.zero_grad()
    
    def get_lr(self):
        return self.optimizer.param_groups[0]['lr']


def build_codec_model(config, device):
    """Build mel codec model for validation (decoding predicted tokens)."""
    model_params = config.get('model_params', {})
    decoder_params = model_params.get('decoder', {})
    
    codec = MelCodecVocoderSpeaker(
        num_speakers=model_params.get('num_speakers', model_params.get('n_speakers', 11)),
        speaker_dim=model_params.get('speaker_dim', 128),
        n_mels=model_params.get('n_mels', 80),
        latent_dim=model_params.get('latent_dim', 512),
        hidden_dim=model_params.get('hidden_dim', 512),
        codec_strides=model_params.get('codec_strides', [2, 2]),
        fsq_levels=model_params.get('fsq_levels', [4] * 6),
        upsample_rates=decoder_params.get('upsample_rates', model_params.get('upsample_rates', [7, 9])),
        gen_istft_n_fft=decoder_params.get('gen_istft_n_fft', model_params.get('gen_istft_n_fft', 28)),
        gen_istft_hop_size=decoder_params.get('gen_istft_hop_size', model_params.get('gen_istft_hop_size', 7)),
        sample_rate=44100,
        n_phonemes=model_params.get('n_phonemes', 178),
        num_languages=model_params.get('num_languages', 10),
        language_dim=model_params.get('language_dim', 64),
    )
    
    return codec.to(device)


@click.command()
@click.option('-p', '--config_path', default='Configs/config_fsq_flow_convnext.yml', type=str)
@click.option('--codec_config_path', default='Configs/config_codec_mel_speaker.yml', type=str)
@click.option('--pretrained_flow_path', default='', type=str, help="Path to pretrained flow predictor checkpoint")
def main(config_path: str, codec_config_path: str, pretrained_flow_path: str):
    # Load configs
    config = yaml.safe_load(open(config_path, "r"))
    codec_config = yaml.safe_load(open(codec_config_path, "r")) if osp.exists(codec_config_path) else {}
    
    # Override log dir for duration training to avoid overwriting flow training
    log_dir = config["log_dir"] + "_duration"
    os.makedirs(log_dir, exist_ok=True)
    shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path)))
    
    # Accelerator setup
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(
        project_dir=log_dir,
        split_batches=True,
        kwargs_handlers=[ddp_kwargs],
        mixed_precision="bf16",
    )
    device = accelerator.device
    
    if accelerator.is_main_process:
        writer = SummaryWriter(osp.join(log_dir, "tensorboard"))
    
    # Logging
    if accelerator.is_main_process:
        py_logger = logger.logger
        py_logger.setLevel(logging.DEBUG)
        py_logger.propagate = False
        formatter = logging.Formatter("%(levelname)s:%(asctime)s: %(message)s")
        if not py_logger.handlers:
            console_handler = logging.StreamHandler(sys.stdout)
            console_handler.setLevel(logging.INFO)
            console_handler.setFormatter(formatter)
            py_logger.addHandler(console_handler)
            
            log_file_path = osp.abspath(osp.join(log_dir, "train.log"))
            file_handler = logging.FileHandler(log_file_path)
            file_handler.setLevel(logging.DEBUG)
            file_handler.setFormatter(formatter)
            py_logger.addHandler(file_handler)
    
    # ==========================================================================
    # CONFIG PARAMETERS
    # ==========================================================================
    batch_size = int(config.get("batch_size", 16))
    epochs = int(config.get("epochs", 100))
    save_freq = int(config.get("save_freq", 2))
    log_interval = int(config.get("log_interval", 20))
    val_interval = int(config.get("val_interval", 50000))
    
    # Data params
    data_params = config.get("data_params", {})
    train_path = data_params.get("train_data", "Data/train_list.txt")
    val_path = data_params.get("val_data", "Data/val_list.txt")
    min_length = int(data_params.get("min_length", 50))
    n_mels = int(data_params.get("n_mels", 80))
    sample_rate = int(config.get("sample_rate", 44100))
    
    # Model params
    model_params = config.get("model_params", {})
    vocab_size = int(model_params.get("vocab_size", 178))
    text_dim = int(model_params.get("text_dim", 512))
    text_num_layers = int(model_params.get("text_num_layers", 6))
    n_speakers = int(model_params.get("n_speakers", 11))
    speaker_dim = int(model_params.get("speaker_dim", 128))
    num_languages = int(model_params.get("num_languages", 10))
    language_dim = int(model_params.get("language_dim", 64))
    hidden_dim = int(model_params.get("hidden_dim", 512))
    flow_num_layers = int(model_params.get("flow_num_layers", 8))
    flow_num_heads = int(model_params.get("flow_num_heads", 8))
    duration_n_flows = int(model_params.get("duration_n_flows", 4))
    fsq_levels = model_params.get("fsq_levels", [4, 4, 4, 4, 4, 4])
    codec_stride = int(model_params.get("codec_stride", 2))
    n_phonemes = int(model_params.get("n_phonemes", 178))
    
    # Training params
    lr = float(config.get("lr", 1e-5))
    lr_decay_steps = int(config.get("lr_decay_steps", 800000))
    lr_final = float(config.get("lr_final", 2e-5))
    duration_pred_weight = float(config.get("duration_pred_weight", 1.0)) # Focus on this
    
    # Paths
    alignment_path_train = config.get("alignment_train", "../data_preparation_v2/alignments_train.safetensors")
    alignment_path_val = config.get("alignment_val", "../data_preparation_v2/alignments_val.safetensors")
    z_continuous_dir_train = config.get("z_continuous_train", "precomputed_tokens/train/z_continuous")
    z_continuous_dir_val = config.get("z_continuous_val", "precomputed_tokens/val/z_continuous")
    
    # ==========================================================================
    # DATA LOADING
    # ==========================================================================
    train_list, val_list = get_data_path_list(train_path, val_path)
    
    train_loader = build_fsq_flow_dataloader(
        train_list,
        z_continuous_dir=z_continuous_dir_train,
        alignment_path=alignment_path_train,
        validation=False,
        min_length=min_length,
        batch_size=batch_size,
        num_workers=4,
        device=device,
        balance_speakers=True,
    )
    
    val_loader = build_fsq_flow_dataloader(
        val_list,
        z_continuous_dir=z_continuous_dir_val,
        alignment_path=alignment_path_val,
        validation=True,
        min_length=min_length,
        batch_size=1,
        num_workers=0,
        device=device,
        balance_speakers=False,
    )
    
    # ==========================================================================
    # BUILD MODEL
    # ==========================================================================
    logger.info("Building models...")

    # 1. Flow Predictor (built first so we can share its ConvNeXt text encoder)
    flow_predictor = FSQFlowPredictorConvNeXt(
        vocab_size=vocab_size,
        text_dim=text_dim,
        text_num_layers=text_num_layers,
        text_kernel_size=7,
        num_languages=num_languages,
        language_dim=language_dim,
        n_speakers=n_speakers,
        speaker_dim=speaker_dim,
        fsq_levels=fsq_levels,
        hidden_dim=hidden_dim,
        cond_layers=flow_num_layers,
        flow_num_heads=flow_num_heads,
        n_phonemes=n_phonemes,
        dropout=0.1,
        codec_stride=codec_stride,
        pad_token_id=0,
    )

    # Load pretrained flow predictor
    if pretrained_flow_path and osp.exists(pretrained_flow_path):
        logger.info(f"Loading pretrained flow predictor from {pretrained_flow_path}")
        state = torch.load(pretrained_flow_path, map_location='cpu')
        if "predictor" in state:
            flow_predictor.load_state_dict(state["predictor"], strict=False)
        else:
            flow_predictor.load_state_dict(state, strict=False)
    else:
        logger.warning(f"No pretrained flow predictor path provided or found at {pretrained_flow_path}. Validation audio will be random noise.")

    # Freeze entire flow predictor, then unfreeze only the text encoder so it
    # receives gradients from the duration loss (joint training of shared encoder).
    flow_predictor.eval()
    for p in flow_predictor.parameters():
        p.requires_grad = False
    for p in flow_predictor.text_encoder.parameters():
        p.requires_grad = True
    flow_predictor.to(device)

    # 2. Duration Predictor (TRAINABLE) — shares ConvNeXt text encoder with flow predictor
    duration_predictor = PredictorDurationSpeakerContext(
        text_encoder=flow_predictor.text_encoder,
        vocab_size=vocab_size,
        n_speakers=n_speakers,
        speaker_emb_dim=speaker_dim,
        filter_channels=hidden_dim,
        p_dropout=0.1,
        n_flows=duration_n_flows,
        pad_token_id=0,
    )

    # 3. Codec (FROZEN - for validation)
    codec = None
    codec_ckpt = config.get("pretrained_codec", "")
    if codec_ckpt and osp.exists(codec_ckpt):
        print(f"Loading codec from {codec_ckpt}")
        codec = build_codec_model(codec_config, device)
        state = torch.load(codec_ckpt, map_location='cpu', weights_only=False)
        if 'net' in state:
            codec.load_state_dict(state['net'].get('codec', state['net']), strict=False)
        else:
            codec.load_state_dict(state, strict=False)
        codec.eval()
        for p in codec.parameters():
            p.requires_grad = False
    
    # Optimizer — covers duration predictor head + shared text encoder
    # (text encoder params are included because it is registered as duration_predictor.text_encoder)
    optimizer = Optimizer(
        list(duration_predictor.parameters()),
        lr=lr,
    )
    
    # Prepare with accelerator
    duration_predictor = accelerator.prepare(duration_predictor)
    optimizer.optimizer = accelerator.prepare(optimizer.optimizer)
    
    # Flow and codec are already on device (if not wrapped) but let's make sure
    # (Since they are frozen, we don't need to wrap them with DDP usually unless we need gradients, which we don't)
    # However, for inference in validation, being on the right device matters.
    
    duration_predictor = duration_predictor.to(device)
    
    optimizer.set_scheduler(total_steps=lr_decay_steps, start_lr=lr, final_lr=lr_final)
    if optimizer.scheduler is not None:
        optimizer.scheduler = accelerator.prepare(optimizer.scheduler)
        
    # Load checkpoint if exists
    start_epoch = 0
    iters = 0
    best_val_loss = float('inf')
    
    # Try to resume duration predictor training
    potential_ckpt = osp.join(log_dir, "latest_model.pt")
    if osp.exists(potential_ckpt):
        logger.info(f"Resuming from {potential_ckpt}")
        state = torch.load(potential_ckpt, map_location='cpu')
        if "duration_predictor" in state:
            duration_predictor.load_state_dict(state["duration_predictor"])
        if "optimizer" in state:
            optimizer.optimizer.load_state_dict(state["optimizer"])
        start_epoch = state.get("epoch", 0)
        iters = state.get("iters", 0)
        best_val_loss = state.get("best_val_loss", float('inf'))
    
    # ==========================================================================
    # TRAINING LOOP
    # ==========================================================================
    logger.info(f"Starting training from epoch {start_epoch}...")
    
    for epoch in range(start_epoch, epochs):
        duration_predictor.train()
        
        epoch_loss = 0.0
        num_batches = 0
        
        for batch_idx, batch in enumerate(train_loader):
            optimizer.zero_grad()
            
            # Unpack batch
            speakers = batch[1].to(device)
            langs = batch[2].to(device)
            # We don't need fsq/codec stuff for training duration
            tgt_texts = batch[5].to(device)          # [B, T_text]
            tgt_lens = batch[6].to(device)
            ctx_texts = batch[7].to(device)
            ctx_lens = batch[8].to(device)
            durations = batch[9].to(device)          # [B, T_dur]
            token_start_indices = batch[10].to(device)
            token_end_indices = batch[11].to(device)
            
            phonemes = tgt_texts
            phoneme_lengths = tgt_lens
            dur_for_model = durations

            dur_output = duration_predictor(
                x_tokens=phonemes,
                x_lengths=phoneme_lengths,
                w_durations=dur_for_model,
                speaker_ids=speakers,
                language_ids=langs,
                ctx_tokens=ctx_texts,
                ctx_lens=ctx_lens,
                token_start_indices=token_start_indices,
                token_end_indices=token_end_indices,
            )
            
            loss_dur_pred = dur_output['loss_dur']
            loss_total = duration_pred_weight * loss_dur_pred
            
            accelerator.backward(loss_total)
            optimizer.step()
            
            epoch_loss += loss_dur_pred.item()
            num_batches += 1
            iters += 1
            
            if accelerator.is_main_process and batch_idx % log_interval == 0:
                writer.add_scalar("train/loss_dur_pred", loss_dur_pred.item(), iters)
                writer.add_scalar("train/lr", optimizer.get_lr(), iters)
                logger.info(
                    f"Epoch [{epoch+1}/{epochs}] Batch [{batch_idx}/{len(train_loader)}] "
                    f"Loss: {loss_dur_pred.item():.4f} LR: {optimizer.get_lr():.6f}"
                )
            
            # Validation
            if accelerator.is_main_process and iters % val_interval == 0:
                val_loss = validate(
                    duration_predictor, flow_predictor, val_loader,
                    codec_stride, device, writer, iters, codec,
                    duration_pred_weight=duration_pred_weight,
                    sample_rate=sample_rate, num_audio_samples=10
                )
                
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    save_checkpoint(
                        duration_predictor, optimizer, epoch, iters, best_val_loss,
                        osp.join(log_dir, "best_model.pt")
                    )
                    logger.info(f"New best model saved! Val Loss: {val_loss:.4f}")
                
                duration_predictor.train()
        
        # Save epoch checkpoint
        if accelerator.is_main_process and (epoch + 1) % save_freq == 0:
            save_checkpoint(
                duration_predictor, optimizer, epoch, iters, best_val_loss,
                osp.join(log_dir, f"model_epoch_{epoch+1}.pt")
            )
            save_checkpoint(
                duration_predictor, optimizer, epoch, iters, best_val_loss,
                osp.join(log_dir, "latest_model.pt")
            )

    if accelerator.is_main_process:
        writer.close()
        logger.info("Training complete!")

def validate(
    duration_predictor,
    flow_predictor,
    val_loader,
    codec_stride,
    device,
    writer,
    iters,
    codec=None,
    duration_pred_weight=1.0,
    sample_rate=44100,
    num_audio_samples=5,
):
    """Run validation and full OOD inference."""
    duration_predictor.eval()
    
    total_loss_dur_pred = 0.0
    num_samples = 0
    audio_samples_generated = 0
    
    with torch.no_grad():
        for batch_idx, batch in enumerate(val_loader):
            if batch_idx >= 50: break
            
            speakers = batch[1].to(device)
            langs = batch[2].to(device)
            tgt_texts = batch[5].to(device)
            tgt_lens = batch[6].to(device)
            ctx_texts = batch[7].to(device)
            ctx_lens = batch[8].to(device)
            durations = batch[9].to(device)
            token_start_indices = batch[10].to(device)
            token_end_indices = batch[11].to(device)
            
            phonemes = tgt_texts
            phoneme_lengths = tgt_lens
            dur_for_model = durations
            
            dur_output = duration_predictor(
                x_tokens=phonemes,
                x_lengths=phoneme_lengths,
                w_durations=dur_for_model,
                speaker_ids=speakers,
                language_ids=langs,
                ctx_tokens=ctx_texts,
                ctx_lens=ctx_lens,
                token_start_indices=token_start_indices,
                token_end_indices=token_end_indices,
            )

            loss_dur_pred = dur_output['loss_dur']
            total_loss_dur_pred += loss_dur_pred.item()
            num_samples += 1
            
            # GENERATE AUDIO (Full OOD Pipeline)
            if codec is not None and flow_predictor is not None and audio_samples_generated < num_audio_samples:
                try:
                    B = speakers.shape[0]
                    unwrapped_dur_pred = getattr(duration_predictor, 'module', duration_predictor)
                    
                    # 1. Predict Duration
                    pred_dur_ood, _, _ = unwrapped_dur_pred.infer(
                        x_tokens=phonemes,
                        x_lengths=phoneme_lengths,
                        speaker_ids=speakers,
                        language_ids=langs,
                        ctx_tokens=ctx_texts,
                        ctx_lens=ctx_lens,
                        token_start_indices=token_start_indices,
                        token_end_indices=token_end_indices,
                    )
                    
                    # 2. Predict Flow Embeddings (using predicted duration)
                    inference_out_ood = flow_predictor.inference(
                        phonemes=phonemes,
                        phoneme_lengths=phoneme_lengths,
                        durations=pred_dur_ood,
                        speaker_ids=speakers,
                        language_ids=langs,
                        n_steps=20,
                        temperature=0.6,
                        ctx_tokens=ctx_texts,
                        ctx_lens=ctx_lens,
                        token_start_indices=token_start_indices,
                        token_end_indices=token_end_indices,
                    )
                    
                    pred_embeddings_ood = inference_out_ood['pred_embeddings']

                    for b in range(min(B, num_audio_samples - audio_samples_generated)):
                        speaker_id = speakers[b:b+1]
                        lang_id = langs[b:b+1]
                        
                        # Get lengths
                        frame_len_ood = pred_dur_ood[b].sum().long().item()
                        codec_len_ood = max(1, (frame_len_ood + codec_stride - 1) // codec_stride)
                        
                        z_pred_ood = pred_embeddings_ood[b:b+1, :codec_len_ood]
                        
                        # 3. Decode Audio
                        wav_pred_ood, _, _ = codec.decode_from_continuous(
                            z_pred_ood, speaker_id, lang_id, f0=None
                        )
                        
                        writer.add_audio(
                            f'val/ood_sample_{audio_samples_generated}_spk{speaker_id.item()}',
                            wav_pred_ood.cpu().numpy().squeeze(),
                            iters,
                            sample_rate=sample_rate
                        )
                        
                        audio_samples_generated += 1
                        
                except Exception as e:
                    logger.warning(f"Audio gen failed: {e}")
                    
    avg_loss_dur = total_loss_dur_pred / max(num_samples, 1)
    writer.add_scalar("val/loss_dur_pred", avg_loss_dur, iters)
    logger.info(f"Validation Loss (Duration): {avg_loss_dur:.4f}")
    return avg_loss_dur

def save_checkpoint(duration_predictor, optimizer, epoch, iters, best_val_loss, path):
    state = {
        "duration_predictor": duration_predictor.state_dict(),
        "optimizer": optimizer.optimizer.state_dict(),
        "epoch": epoch + 1,
        "iters": iters,
        "best_val_loss": best_val_loss,
    }
    torch.save(state, path)
    logger.info(f"Checkpoint saved to {path}")

if __name__ == "__main__":
    main()