from pathlib import Path

import soundfile as sf
import torch
import hydra
from accelerate import Accelerator

from omegaconf import OmegaConf
from safetensors.torch import load_file
import diffusers.schedulers as noise_schedulers
from tqdm import tqdm

from uniflow_audio.utils.config import register_omegaconf_resolvers
from uniflow_audio.models.common import LoadPretrainedBase
from uniflow_audio.utils.general import sanitize_filename

try:
    import torch_npu
    from torch_npu.contrib import transfer_to_npu
except:
    pass

register_omegaconf_resolvers()


def main():

    accelerator = Accelerator(mixed_precision="no")
    configs = []

    @hydra.main(config_path="configs", config_name="inference")
    def parse_config_from_command_line(config):
        config = OmegaConf.to_container(config, resolve=True)
        configs.append(config)

    parse_config_from_command_line()
    config = configs[0]

    if "exp_dir" in config:
        use_best = config.get("use_best", True)
        exp_dir = Path(config["exp_dir"])
        if use_best:  # use best ckpt
            ckpt_path: Path = sorted((exp_dir / "checkpoints").iterdir()
                                    )[0] / "model.safetensors"
        else:  # use last ckpt
            ckpt_path: Path = sorted((exp_dir / "checkpoints").iterdir()
                                    )[-1] / "model.safetensors"
    elif "ckpt_dir" in config:
        ckpt_dir = Path(config["ckpt_dir"])
        ckpt_path = ckpt_dir / "model.safetensors"
        exp_dir = ckpt_dir.parent.parent

    accelerator.print(f'\n ckpt path: {ckpt_path}\n ')

    exp_config = OmegaConf.load(exp_dir / "config.yaml")
    model: LoadPretrainedBase = hydra.utils.instantiate(exp_config["model"])
    state_dict = load_file(ckpt_path)
    model.load_pretrained(state_dict)

    if "sampler" in config["test_dataloader"]:
        data_source = hydra.utils.instantiate(
            config["test_dataloader"]["dataset"], _convert_="all"
        )
        sampler = hydra.utils.instantiate(
            config["test_dataloader"]["sampler"],
            data_source=data_source,
            _convert_="all"
        )
        test_dataloader = hydra.utils.instantiate(
            config["test_dataloader"], sampler=sampler, _convert_="all"
        )
    else:
        test_dataloader = hydra.utils.instantiate(
            config["test_dataloader"], _convert_="all"
        )

    model.eval()

    model, test_dataloader = accelerator.prepare(model, test_dataloader)

    scheduler = getattr(
        noise_schedulers,
        config["noise_scheduler"]["type"],
    ).from_pretrained(
        config["noise_scheduler"]["name"],
        subfolder="scheduler",
    )

    audio_output_dir = exp_dir / config["wav_dir"]
    if accelerator.is_main_process:
        audio_output_dir.mkdir(parents=True, exist_ok=True)
    unwrapped_model = accelerator.unwrap_model(model)
    pbar_disable = not accelerator.is_main_process
    cfg = f"cfg_{config['infer_args']['guidance_scale']}"
    step = f"CLB_mixed_step_{config['infer_args']['num_steps']}"
    with torch.no_grad():
        for batch in tqdm(test_dataloader, disable=pbar_disable):
            kwargs = config["infer_args"].copy()
            kwargs.update(batch)

            waveform = unwrapped_model.inference(
                scheduler=scheduler,
                **kwargs,
            )

            for name, wave, task in zip(
                batch["item_name"], waveform, batch["task"]
            ):
                (audio_output_dir / task / cfg / step).mkdir(parents=True, exist_ok=True)
                safe_name = sanitize_filename(name)
                sf.write(
                    audio_output_dir / task / cfg / step / f"{safe_name}.wav",
                    wave[0].cpu().numpy(),
                    samplerate=exp_config["sample_rate"],
                )

    accelerator.wait_for_everyone()


if __name__ == "__main__":
    main()
