# Anima LECO LoRA training script

import argparse
import importlib
import random
from typing import List, Sequence

import torch
from accelerate.utils import set_seed
from tqdm import tqdm

from library.device_utils import clean_memory_on_device, init_ipex

init_ipex()

import train_network
from anima_train_network import AnimaNetworkTrainer
from library import anima_models, anima_train_utils, sd3_train_utils, strategy_base, train_util
from library.leco_train_util import PromptEmbedsCache, build_network_kwargs, get_random_resolution, load_prompt_settings, save_weights
from library.utils import setup_logging

setup_logging()
import logging

logger = logging.getLogger(__name__)


AnimaPromptConds = List[torch.Tensor]


def setup_parser() -> argparse.ArgumentParser:
    parser = train_network.setup_parser()
    train_util.add_dit_training_arguments(parser)
    anima_train_utils.add_anima_training_arguments(parser)

    parser.add_argument("--prompts_file", type=str, required=True, help="LECO prompt toml")
    parser.add_argument(
        "--max_denoising_steps",
        type=int,
        default=40,
        help="number of partial denoising steps per iteration",
    )
    parser.add_argument(
        "--leco_denoise_guidance_scale",
        type=float,
        default=3.0,
        help="guidance scale for the partial denoising pass",
    )
    parser.add_argument(
        "--unsloth_offload_checkpointing",
        action="store_true",
        help="offload activations to CPU RAM using async non-blocking transfers. "
        "Cannot be used with --cpu_offload_checkpointing or --blocks_to_swap.",
    )

    return parser


def encode_prompt_anima(tokenize_strategy, text_encoding_strategy, text_encoders, prompt: str) -> AnimaPromptConds:
    tokens = tokenize_strategy.tokenize(prompt)
    encoded = text_encoding_strategy.encode_tokens(tokenize_strategy, text_encoders, tokens)
    return [tensor.detach().to("cpu") for tensor in encoded]


def _repeat_cond(cond: torch.Tensor, batch_size: int, device: torch.device, dtype: torch.dtype | None = None) -> torch.Tensor:
    if dtype is not None and cond.dtype.is_floating_point:
        cond = cond.to(device=device, dtype=dtype)
    else:
        cond = cond.to(device=device)
    return cond.repeat_interleave(batch_size, dim=0)


def concat_prompt_conds_anima(
    unconditional: Sequence[torch.Tensor],
    conditional: Sequence[torch.Tensor],
    batch_size: int,
    device: torch.device,
    weight_dtype: torch.dtype,
) -> AnimaPromptConds:
    prompt_embeds = torch.cat(
        [
            _repeat_cond(unconditional[0], batch_size, device, weight_dtype),
            _repeat_cond(conditional[0], batch_size, device, weight_dtype),
        ],
        dim=0,
    )
    attn_mask = torch.cat([_repeat_cond(unconditional[1], batch_size, device), _repeat_cond(conditional[1], batch_size, device)], dim=0)
    t5_input_ids = torch.cat([_repeat_cond(unconditional[2], batch_size, device), _repeat_cond(conditional[2], batch_size, device)], dim=0).long()
    t5_attn_mask = torch.cat([_repeat_cond(unconditional[3], batch_size, device), _repeat_cond(conditional[3], batch_size, device)], dim=0)
    return [prompt_embeds, attn_mask, t5_input_ids, t5_attn_mask]


def apply_noise_offset_anima(latents: torch.Tensor, noise_offset: float | None) -> torch.Tensor:
    if noise_offset is None or noise_offset == 0:
        return latents
    noise = torch.randn(
        (latents.shape[0], latents.shape[1], 1, 1, 1),
        dtype=torch.float32,
        device="cpu",
    ).to(dtype=latents.dtype, device=latents.device)
    return latents + noise_offset * noise


def get_initial_latents_anima(batch_size: int, height: int, width: int) -> torch.Tensor:
    latent_h = height // 8
    latent_w = width // 8
    return torch.randn((batch_size, 16, 1, latent_h, latent_w), device="cpu")


def predict_anima(
    anima: anima_models.Anima,
    timestep,
    latents: torch.Tensor,
    prompt_conds: AnimaPromptConds,
    guidance_scale: float = 1.0,
) -> torch.Tensor:
    batch_size = latents.shape[0]
    model_input = torch.cat([latents, latents], dim=0)

    if not torch.is_tensor(timestep):
        timestep = torch.tensor(timestep, device=latents.device, dtype=latents.dtype)
    timestep = timestep.to(device=latents.device, dtype=latents.dtype)
    if timestep.ndim == 0:
        timestep = timestep.expand(model_input.shape[0])
    elif timestep.shape[0] == batch_size:
        timestep = timestep.repeat(2)

    padding_mask = torch.zeros(
        model_input.shape[0],
        1,
        latents.shape[-2],
        latents.shape[-1],
        dtype=latents.dtype,
        device=latents.device,
    )

    if hasattr(anima, "prepare_block_swap_before_forward"):
        anima.prepare_block_swap_before_forward()

    prompt_embeds, attn_mask, t5_input_ids, t5_attn_mask = prompt_conds
    model_pred = anima(
        model_input,
        timestep,
        prompt_embeds,
        padding_mask=padding_mask,
        target_input_ids=t5_input_ids,
        target_attention_mask=t5_attn_mask,
        source_attention_mask=attn_mask,
    )
    pred_uncond, pred_text = model_pred.chunk(2, dim=0)
    return pred_uncond + guidance_scale * (pred_text - pred_uncond)


def diffusion_anima(
    anima: anima_models.Anima,
    noise_scheduler,
    latents: torch.Tensor,
    prompt_conds: AnimaPromptConds,
    total_timesteps: int,
    guidance_scale: float = 3.0,
) -> torch.Tensor:
    sigmas = noise_scheduler.sigmas.to(device=latents.device, dtype=latents.dtype)
    for step_index in range(total_timesteps):
        sigma = sigmas[step_index]
        model_pred = predict_anima(anima, sigma, latents, prompt_conds, guidance_scale=guidance_scale)
        dt = sigmas[step_index + 1] - sigma
        latents = (latents + model_pred * dt).to(latents.dtype)
    return latents


def main():
    parser = setup_parser()
    args = parser.parse_args()
    args = train_util.read_config_from_file(args, parser)
    if args.network_module is None:
        args.network_module = "networks.lora_anima"
    if args.network_dim is None:
        args.network_dim = 4
    train_util.verify_training_args(args)
    setup_logging(args, reset=True)

    if args.output_dir is None:
        raise ValueError("--output_dir is required")
    if args.network_train_text_encoder_only:
        raise ValueError("Anima LECO does not support text encoder LoRA training")
    args.network_train_unet_only = True
    if args.fp8_base or args.fp8_base_unet:
        logger.warning("fp8_base and fp8_base_unet are not supported for Anima LECO. Disabling them.")
        args.fp8_base = False
        args.fp8_base_unet = False
    args.fp8_scaled = False

    if args.seed is None:
        args.seed = random.randint(0, 2**32 - 1)
    set_seed(args.seed)

    trainer = AnimaNetworkTrainer()
    accelerator = train_util.prepare_accelerator(args)
    weight_dtype, save_dtype = train_util.prepare_dtype(args)

    prompt_settings = load_prompt_settings(args.prompts_file)
    logger.info(f"loaded {len(prompt_settings)} LECO prompt settings from {args.prompts_file}")

    tokenize_strategy = trainer.get_tokenize_strategy(args)
    text_encoding_strategy = trainer.get_text_encoding_strategy(args)
    strategy_base.TokenizeStrategy.set_strategy(tokenize_strategy)
    strategy_base.TextEncodingStrategy.set_strategy(text_encoding_strategy)

    _, text_encoders, vae, unet = trainer.load_target_model(args, weight_dtype, accelerator)
    del vae
    text_encoders[0].to(accelerator.device, dtype=weight_dtype)
    text_encoders[0].requires_grad_(False)
    text_encoders[0].eval()

    prompt_cache = PromptEmbedsCache()
    unique_prompts = sorted(
        {
            prompt
            for setting in prompt_settings
            for prompt in (setting.target, setting.positive, setting.unconditional, setting.neutral)
        }
    )
    with torch.no_grad(), accelerator.autocast():
        for prompt in unique_prompts:
            prompt_cache[prompt] = encode_prompt_anima(tokenize_strategy, text_encoding_strategy, text_encoders, prompt)

    text_encoders[0].to("cpu")
    clean_memory_on_device(accelerator.device)

    unet, text_encoders = trainer.load_unet_lazily(args, weight_dtype, accelerator, text_encoders)
    unet.requires_grad_(False)
    unet.train()

    network_module = importlib.import_module(args.network_module)
    net_kwargs = build_network_kwargs(args)
    if args.dim_from_weights:
        if args.network_weights is None:
            raise ValueError("--dim_from_weights requires --network_weights")
        network, _ = network_module.create_network_from_weights(1.0, args.network_weights, None, text_encoders, unet, **net_kwargs)
    else:
        network = network_module.create_network(
            1.0,
            args.network_dim,
            args.network_alpha,
            None,
            text_encoders,
            unet,
            neuron_dropout=args.network_dropout,
            **net_kwargs,
        )

    network.apply_to(text_encoders, unet, apply_text_encoder=False, apply_unet=True)
    network.set_multiplier(0.0)

    if args.network_weights is not None:
        info = network.load_weights(args.network_weights)
        logger.info(f"loaded network weights from {args.network_weights}: {info}")

    if args.gradient_checkpointing:
        if args.cpu_offload_checkpointing:
            unet.enable_gradient_checkpointing(cpu_offload=True)
        else:
            unet.enable_gradient_checkpointing()
        if args.unsloth_offload_checkpointing:
            unet.enable_gradient_checkpointing(unsloth_offload=True)
        network.enable_gradient_checkpointing()

    unet = trainer.prepare_unet_with_accelerator(args, accelerator, unet)

    unet_lr = args.unet_lr if args.unet_lr is not None else args.learning_rate
    if hasattr(network, "prepare_optimizer_params_with_multiple_te_lrs"):
        trainable_params, _ = network.prepare_optimizer_params_with_multiple_te_lrs(None, unet_lr, args.learning_rate)
    else:
        trainable_params, _ = network.prepare_optimizer_params(None, unet_lr, args.learning_rate)
    _, _, optimizer = train_util.get_optimizer(args, trainable_params)
    lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)

    network, optimizer, lr_scheduler = accelerator.prepare(network, optimizer, lr_scheduler)
    accelerator.unwrap_model(network).prepare_grad_etc(text_encoders, unet)

    if args.full_fp16:
        train_util.patch_accelerator_for_fp16_training(accelerator)

    optimizer_train_fn, _ = train_util.get_optimizer_train_eval_fn(optimizer, args)
    optimizer_train_fn()
    train_util.init_trackers(accelerator, args, "anima_leco_train")

    noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(
        num_train_timesteps=1000,
        shift=args.discrete_flow_shift,
    )

    progress_bar = tqdm(total=args.max_train_steps, disable=not accelerator.is_local_main_process, desc="steps")
    global_step = 0

    while global_step < args.max_train_steps:
        with accelerator.accumulate(network):
            optimizer.zero_grad(set_to_none=True)

            setting = prompt_settings[torch.randint(0, len(prompt_settings), (1,)).item()]
            noise_scheduler.set_timesteps(args.max_denoising_steps, device=accelerator.device)
            timesteps_to = torch.randint(1, args.max_denoising_steps, (1,), device=accelerator.device).item()
            height, width = get_random_resolution(setting)
            height = max(64, height - height % 16)
            width = max(64, width - width % 16)

            latents = get_initial_latents_anima(setting.batch_size, height, width).to(accelerator.device, dtype=weight_dtype)
            latents = apply_noise_offset_anima(latents, args.noise_offset)

            network_multiplier = accelerator.unwrap_model(network)
            network_multiplier.set_multiplier(setting.multiplier)
            with torch.no_grad(), accelerator.autocast():
                denoised_latents = diffusion_anima(
                    unet,
                    noise_scheduler,
                    latents,
                    concat_prompt_conds_anima(
                        prompt_cache[setting.unconditional],
                        prompt_cache[setting.target],
                        setting.batch_size,
                        accelerator.device,
                        weight_dtype,
                    ),
                    total_timesteps=timesteps_to,
                    guidance_scale=args.leco_denoise_guidance_scale,
                )
            denoised_latents = denoised_latents.detach()

            current_sigma = noise_scheduler.sigmas[timesteps_to].to(device=accelerator.device, dtype=weight_dtype)

            network_multiplier.set_multiplier(0.0)
            with torch.no_grad(), accelerator.autocast():
                positive_latents = predict_anima(
                    unet,
                    current_sigma,
                    denoised_latents,
                    concat_prompt_conds_anima(
                        prompt_cache[setting.unconditional],
                        prompt_cache[setting.positive],
                        setting.batch_size,
                        accelerator.device,
                        weight_dtype,
                    ),
                    guidance_scale=1.0,
                )
                neutral_latents = predict_anima(
                    unet,
                    current_sigma,
                    denoised_latents,
                    concat_prompt_conds_anima(
                        prompt_cache[setting.unconditional],
                        prompt_cache[setting.neutral],
                        setting.batch_size,
                        accelerator.device,
                        weight_dtype,
                    ),
                    guidance_scale=1.0,
                )
                unconditional_latents = predict_anima(
                    unet,
                    current_sigma,
                    denoised_latents,
                    concat_prompt_conds_anima(
                        prompt_cache[setting.unconditional],
                        prompt_cache[setting.unconditional],
                        setting.batch_size,
                        accelerator.device,
                        weight_dtype,
                    ),
                    guidance_scale=1.0,
                )

            network_multiplier.set_multiplier(setting.multiplier)
            with accelerator.autocast():
                target_latents = predict_anima(
                    unet,
                    current_sigma,
                    denoised_latents,
                    concat_prompt_conds_anima(
                        prompt_cache[setting.unconditional],
                        prompt_cache[setting.target],
                        setting.batch_size,
                        accelerator.device,
                        weight_dtype,
                    ),
                    guidance_scale=1.0,
                )

                target = setting.build_target(positive_latents, neutral_latents, unconditional_latents)
                loss = torch.nn.functional.mse_loss(target_latents.float(), target.float(), reduction="none")
                loss = loss.mean(dim=list(range(1, loss.ndim)))
                sigma_for_weighting = current_sigma.expand(loss.shape[0]).view(-1, 1, 1, 1, 1)
                weighting = anima_train_utils.compute_loss_weighting_for_anima(args.weighting_scheme, sigma_for_weighting)
                loss = loss * weighting.view(loss.shape[0])
                loss = loss.mean() * setting.weight

            accelerator.backward(loss)

            if accelerator.sync_gradients and args.max_grad_norm != 0.0:
                accelerator.clip_grad_norm_(network.parameters(), args.max_grad_norm)

            optimizer.step()
            lr_scheduler.step()

        if accelerator.sync_gradients:
            global_step += 1
            progress_bar.update(1)
            accelerator.unwrap_model(network).set_multiplier(0.0)

            logs = {
                "loss": loss.detach().item(),
                "lr": lr_scheduler.get_last_lr()[0],
                "guidance_scale": setting.guidance_scale,
                "network_multiplier": setting.multiplier,
                "sigma": current_sigma.detach().float().item(),
            }
            accelerator.log(logs, step=global_step)
            progress_bar.set_postfix(loss=f"{logs['loss']:.4f}")

            if args.save_every_n_steps and global_step % args.save_every_n_steps == 0 and global_step < args.max_train_steps:
                accelerator.wait_for_everyone()
                if accelerator.is_main_process:
                    extra_metadata = {"ss_base_model_version": "anima_preview"}
                    save_weights(accelerator, network, args, save_dtype, prompt_settings, global_step, last=False, extra_metadata=extra_metadata)

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        extra_metadata = {"ss_base_model_version": "anima_preview"}
        save_weights(accelerator, network, args, save_dtype, prompt_settings, global_step, last=True, extra_metadata=extra_metadata)

    accelerator.end_training()


if __name__ == "__main__":
    main()
