from pathlib import Path

import hydra
import torch
from omegaconf import DictConfig
from slider import Beatmap

from osudiffusion import DiT_models
from osuT5.inference import Preprocessor, Pipeline, Postprocessor, DiffisionPipeline
from osuT5.tokenizer import Tokenizer
from osuT5.utils import get_model


def get_args_from_beatmap(args: DictConfig):
    if args.beatmap_path is None or args.beatmap_path == "":
        return

    beatmap_path = Path(args.beatmap_path)

    if not beatmap_path.is_file():
        raise FileNotFoundError(f"Beatmap file {beatmap_path} not found.")

    beatmap = Beatmap.from_path(beatmap_path)
    args.audio_path = beatmap_path.parent / beatmap.audio_filename
    args.output_path = beatmap_path.parent
    args.bpm = beatmap.bpm_max()
    args.offset = min(tp.offset.total_seconds() * 1000 for tp in beatmap.timing_points)
    args.slider_multiplier = beatmap.slider_multiplier
    args.title = beatmap.title
    args.artist = beatmap.artist
    args.beatmap_id = beatmap.beatmap_id if args.beatmap_id == -1 else args.beatmap_id
    args.diffusion.style_id = beatmap.beatmap_id if args.diffusion.style_id == -1 else args.diffusion.style_id
    args.difficulty = float(beatmap.stars()) if args.difficulty == -1 else args.difficulty


def find_model(ckpt_path, args: DictConfig, device):
    assert Path(ckpt_path).exists(), f"Could not find DiT checkpoint at {ckpt_path}"
    checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
    if "ema" in checkpoint:  # supports checkpoints from train.py
        checkpoint = checkpoint["ema"]

    model = DiT_models[args.diffusion.model](
        num_classes=args.diffusion.num_classes,
        context_size=19 - 3 + 128,
    ).to(device)
    model.load_state_dict(checkpoint)
    model.eval()  # important!
    return model


@hydra.main(config_path="configs", config_name="inference", version_base="1.1")
def main(args: DictConfig):
    get_args_from_beatmap(args)

    torch.set_grad_enabled(False)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ckpt_path = Path(args.model_path)
    model_state = torch.load(ckpt_path / "pytorch_model.bin", map_location=device)
    tokenizer_state = torch.load(ckpt_path / "custom_checkpoint_0.pkl")

    tokenizer = Tokenizer()
    tokenizer.load_state_dict(tokenizer_state)

    model = get_model(args, tokenizer)
    model.load_state_dict(model_state)
    model.eval()
    model.to(device)
    
    preprocessor = Preprocessor(args)
    audio = preprocessor.load(args.audio_path)
    sequences = preprocessor.segment(audio)
    total_duration_ms = len(audio) / 16000 * 1000
    args.total_duration_ms = total_duration_ms




    
    generated_maps = []
    generated_positions = []
    diffs = []

    
    if args.full_set:
        for i in range(args.set_difficulties):
            diffs.append(3 + i * (7 - 3) / (args.set_difficulties - 1))
            
        print(diffs)
        for diff in diffs:
            print(f"Generating difficulty {diff}")
            args.difficulty = diff
            pipeline = Pipeline(args, tokenizer)
            events = pipeline.generate(model, sequences)
            generated_maps.append(events)
    else:
        pipeline = Pipeline(args, tokenizer)
        events = pipeline.generate(model, sequences)
        generated_maps.append(events)
    
    
        
    if args.generate_positions:
        model = find_model(args.diff_ckpt, args, device)
        refine_model = find_model(args.diff_refine_ckpt, args, device) if len(args.diff_refine_ckpt) > 0 else None
        diffusion_pipeline = DiffisionPipeline(args.diffusion)
        for events in generated_maps:
            events = diffusion_pipeline.generate(model, events, refine_model)
            generated_positions.append(events)
    else:
        generated_positions = generated_maps
        
    postprocessor = Postprocessor(args)
    postprocessor.generate(generated_positions)


if __name__ == "__main__":
    main()