import argparse
import sys
import pathlib
import torch

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', type=str, metavar='PATH', required=True)
    parser.add_argument('--pyannote-path', type=str, metavar='PATH', required=True)
    parser.add_argument('--stats-output', type=str, metavar='PATH', required=False)
    args = parser.parse_args()

    # setup pyannote
    sys.path.insert(0, args.pyannote_path) # "pyannote-Normal")
    from pyannote.audio.applications.config import load_config, load_specs

    # Load models
    model_root = pathlib.Path("models")
    path1 = model_root / "config.yml"  # config file
    path2 = model_root / "specs.yml"   # specification file
    path3 = model_root / "0050.pt"     # model parameter file

    config = load_config(path1)
    spec = load_specs(path2)
    model = config["get_model_from_specs"](spec)
    model.eval()
    model.load_state_dict(torch.load(path3, map_location="cpu"))

    example_input = torch.rand(1, 16000, 1)
    model.reset()
    model.forward(example_input)
    # model.print_stats()
    stats = model.dump_stats()

    model.reset(
        sinc_normalization_means=stats['sincnet']['means'],
        sinc_normalization_sigmas=stats['sincnet']['sigmas'],
        sinc_normalization_counts=stats['sincnet']['counts'],
    )
    model.forward(example_input)
    model.print_stats()

    model.reset(
        sinc_normalization_means=[v.mean() for v in stats['sincnet']['means']],
        sinc_normalization_sigmas=[v.mean() for v in stats['sincnet']['sigmas']],
        sinc_normalization_counts=[v.sum() for v in stats['sincnet']['counts']],
    )
    model.forward(example_input)
    model.print_stats()

    model.reset()
    y0 = model.forward(example_input[:, :10000])
    y1 = model.forward(example_input[:, 10000:20000])

    model.reset()
    z0 = model.forward(example_input[:, :10000])
    model.push_status()
    z1 = model.forward(torch.zeros_like(example_input[:, 10000:20000]))
    model.pop_status()
    z2 = model.forward(example_input[:, 10000:20000])
    assert((y1 == z2).all())

    model.reset()
    for n in range(0, 10000, 160):
        segment = example_input[:, n:n+160]
        model.push_status()
        z1 = model.forward(torch.zeros_like(segment))
        model.pop_status()
        z0 = model.forward(segment)

        # print('min: {:.3f}, max: {:.3f}, mean: {:.3f}, std: {:.3f}, head: {:.3f}, tail: {:.3f}'.format(z1[:, :, 0].min().item(), z1[:, :, 0].max().item(), z1[:, :, 0].mean().item(), z1[:, :, 0].std().item(), z1[0, 0, 0].item(), z1[0, -1, 0].item()))
        # print(z1[:, 0].min().item(), z1[:, 0].max().item(), z1[:, 0].mean().item(), z1[:, 0].std().item(), z1[0, 0].item(), z1[-1, 0].item())
        # print(z1.shape)
        # print(z1[:, :, 0].min(), z1[:, :, 0].max(), z1[:, :, 0].mean(), z1[:, :, 0].std(), z1[0, 0, 0], z1[0, -1, 0])
        print(z0 - z1)

    # print(z1[:, :200, 0])
    # print(z2[:, :200, 0])

    if args.stats_output:
        with open(args.stats_output, 'w') as f:
            print('static std::vector<size_t> STATS_COUNTS = {', file=f)
            print(', '.join(['{:d}'.format(v.item()) for v in stats['sincnet']['counts']]), file=f)
            print('};', file=f)

            print('static std::vector<std::vector<float>> STATS_MEANS = {', file=f)
            for vec in stats['sincnet']['means']:
                print('{', file=f)
                print(', '.join(['{:.3f}'.format(v.float()) for v in vec]), file=f)
                print('},', file=f)
            print('};', file=f)

            print('static std::vector<std::vector<float>> STATS_SIGMAS = {', file=f)
            for vec in stats['sincnet']['sigmas']:
                print('{', file=f)
                print(', '.join(['{:.3f}'.format(v.float()) for v in vec]), file=f)
                print('},', file=f)
            print('};', file=f)

    # JIT to TorchScript
    jitmodel = torch.jit.script(model)
    jitmodel.save(args.output) # "models/pyannote_jit_script.pt")
