import torch
import os
from os.path import join as pjoin
from dataset.motion import MotionData, load_multiple_dataset
from models import create_model, create_conditional_model
from models.architecture import draw_example, get_pyramid_lengths, FullGenerator
from option import TestOptionParser, TrainOptionParser
from fix_contact import fix_contact_on_file
from models.utils import get_layered_mask
from interactive_utils import sliding_window


def load_all_from_path(save_path, device, use_class=False):
    train_parser = TrainOptionParser()
    args = train_parser.load(pjoin(save_path, 'args.txt'))
    args.device = device
    args.save_path = save_path
    device = torch.device(args.device)

    if not args.multiple_sequences:
        motion_data = MotionData(pjoin(args.bvh_prefix, f'{args.bvh_name}.bvh'),
                                 padding=args.skeleton_aware, use_velo=args.use_velo, repr=args.repr,
                                 contact=args.contact, keep_y_pos=args.keep_y_pos)
        multiple_data = [motion_data]
    else:
        multiple_data = load_multiple_dataset(prefix=args.bvh_prefix, name_list=pjoin(args.bvh_prefix, args.bvh_name),
                                              padding=args.skeleton_aware, use_velo=args.use_velo, repr=args.repr,
                                              contact=args.contact, keep_y_pos=args.keep_y_pos,
                                              no_scale=True)
        motion_data = multiple_data[0]

    lengths = []
    min_len = 10000
    for i in range(len(multiple_data)):
        new_length = get_pyramid_lengths(args, len(multiple_data[i]))
        min_len = min(min_len, len(new_length))
        if args.num_stages_limit != -1:
            new_length = new_length[:args.num_stages_limit]
        lengths.append(new_length)

    for i in range(len(multiple_data)):
        lengths[i] = lengths[i][-min_len:]

    gens = []
    for step, length in enumerate(lengths[0]):
        create = create_conditional_model if args.conditional_generator and step < args.num_conditional_generator else create_model
        gen = create(args, motion_data, evaluation=True)
        try:
            gen_sate = torch.load(pjoin(args.save_path, f'gen{step:03d}.pt'), map_location=device)
        except FileNotFoundError:
            gen_sate = torch.load(pjoin(args.save_path, f'gen{step}.pt'), map_location=device)
        gen.load_state_dict(gen_sate)
        gens.append(gen)
    z_star = torch.load(pjoin(args.save_path, 'z_star.pt'), map_location=device)
    amps = torch.load(pjoin(args.save_path, 'amps.pt'), map_location=device)
    if use_class:
        if isinstance(z_star, list):
            z_star = z_star[0]
        if len(amps.shape) != 1:
            amps = amps[0]
        return FullGenerator(args, motion_data, gens, z_star, amps)
    else:
        if len(amps.shape) == 1:
            amps = amps.unsqueeze(0)
        if isinstance(z_star, torch.Tensor) and len(z_star.shape) == 3:
            z_star = z_star.unsqueeze(0)
        return args, multiple_data, gens, z_star, amps, lengths


def write_multires(imgs, prefix, writer, interpolator, full_lengths=None, requires_con_loss=True):
    os.makedirs(prefix, exist_ok=True)
    length = imgs[-1].shape[-1] if full_lengths is None else full_lengths
    res = []
    for step, img in enumerate(imgs):
        full_length = interpolator(img, length)
        writer(pjoin(prefix, f'{step:02d}.bvh'), full_length)
        velo = full_length[:, -6:-3].norm(dim=1)
        res.append(velo)
    if requires_con_loss:
        res = torch.cat(res, dim=0)
        consistency_loss = torch.nn.MSELoss()(res[1], res[0])
        return consistency_loss


def gen_noise(n_channel, length, full_noise, device):
    if full_noise:
        res = torch.randn((1, n_channel, length)).to(device)
    else:
        res = torch.randn((1, 1, length)).repeat(1, n_channel, 1).to(device)
    return res


def main():
    test_parser = TestOptionParser()
    test_args = test_parser.parse_args()

    args, multiple_data, gens, z_stars, amps, lengths = load_all_from_path(test_args.save_path, test_args.device)
    device = torch.device(args.device)
    n_total_levels = len(gens)

    motion_data = multiple_data[0]

    noise_channel = z_stars[0].shape[1] if args.full_noise else 1

    if len(args.path_to_existing):
        ConGen = load_all_from_path(args.path_to_existing, args.device, use_class=True)
        ConGen.output_mask = get_layered_mask(args.conditional_mode, motion_data.n_rot)
        conds_rec = [motion_data.sample(lengths[0][i]) for i in range(args.num_conditional_generator)]
    else:
        ConGen = None
        conds_rec = None

    print('levels:', lengths)
    save_path = pjoin(args.save_path, 'bvh')
    os.makedirs(save_path, exist_ok=True)

    base_id = 0

    # Evaluate with reconstruct noise
    for i in range(len(multiple_data)):
        motion_data = multiple_data[i]
        imgs = draw_example(gens, 'rec', z_stars[i], lengths[i] + [1], amps[i], 1, args, all_img=True, conds=conds_rec,
                            full_noise=args.full_noise)
        real = motion_data.sample(size=len(motion_data), slerp=args.slerp).to(device)
        motion_data.write(pjoin(save_path, f'gt_{i}.bvh'), real)
        motion_data.write(pjoin(save_path, f'rec_{i}.bvh'), imgs[-1])

        if imgs[-1].shape[-1] == real.shape[-1]:
            rec_loss = torch.nn.MSELoss()(imgs[-1], real).detach().cpu().numpy()
            print(f'rec_loss: {rec_loss.item():.07f}')

    generation_mode = 'manip' if test_args.style_transfer or test_args.keyframe_editing else 'random'

    if test_args.style_transfer:
        manip_data = MotionData(f'{test_args.style_transfer}',
                                padding=args.skeleton_aware, use_velo=args.use_velo, repr=args.repr,
                                contact=args.contact, keep_y_pos=args.keep_y_pos,
                                no_scale=False)
        target_len = len(manip_data)
        target_length = get_pyramid_lengths(args, target_len)
        conds = manip_data.sample(target_length[0])
    elif test_args.keyframe_editing:
        manip_data = MotionData(f'{test_args.keyframe_editing}',
                                padding=args.skeleton_aware, use_velo=args.use_velo, repr=args.repr,
                                contact=args.contact, keep_y_pos=args.keep_y_pos,
                                no_scale=True)
        target_len = len(multiple_data[0])                       # Use original length of training data
        target_length = get_pyramid_lengths(args, target_len)
        conds = manip_data.sample(target_length[0])
    elif test_args.conditional_generation:
        manip_data = MotionData(f'{test_args.conditional_generation}',
                                padding=args.skeleton_aware, use_velo=args.use_velo, repr=args.repr,
                                contact=args.contact, keep_y_pos=args.keep_y_pos,
                                no_scale=True)
        target_len = len(manip_data)  # Use original length of training data
        target_length = get_pyramid_lengths(args, target_len)
        conds = [manip_data.sample(l) for l in target_length[:args.num_conditional_generator]]
        conds_full = conds[-1]
        generation_mode = 'cond'
        if not args.conditional_generator:
            raise Exception('Conditional generation only applicable to conditional generators.')
    elif args.conditional_generator:
        "This is a conditional model, but the condition is not given. Then the condition will be sampled from the ConGen model"
        target_len = test_args.target_length
        target_length = get_pyramid_lengths(args, target_len)
        manip_data = None
        conds = ConGen.random_generate(target_length)
        conds_full = conds[-1]
        conds = conds[:args.num_conditional_generator]
    else:
        target_len = test_args.target_length
        target_length = get_pyramid_lengths(args, target_len)
        manip_data = None
        conds = None

    while len(target_length) > n_total_levels:
        target_length = target_length[1:]
    z_length = target_length[0]

    z_target = gen_noise(noise_channel, z_length, args.full_noise, device)
    z_target *= amps[base_id][0]

    amps2 = amps[base_id].clone()
    amps2[1:] = 0

    if not test_args.interactive:
        imgs = draw_example(gens, generation_mode, z_stars[base_id], target_length, amps2, 1, args, all_img=True,
                            conds=conds, full_noise=args.full_noise, given_noise=[z_target])
    else:
        if not args.conditional_generator:
            raise Exception('Interactive mode only applicable to conditional generators.')
        final_res, imgs = sliding_window(gens, z_stars[base_id], amps2, conds_full, args)

    motion_data.write(pjoin(save_path, f'result.bvh'), imgs[-1])
    fix_contact_on_file(save_path, name=f'result')

    if manip_data is not None:
        motion_data.write(pjoin(save_path, 'manipulate_input.bvh'), manip_data.sample())
    elif args.conditional_generator:
        motion_data.write(pjoin(save_path, 'generated_traj.bvh'), conds_full)


if __name__ == '__main__':
    main()