import json
import os
import pickle

import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np

from .config import FLAGS
from .model import Generator


class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self


def mel2wave(mel):
    config_file = "assets/hifigan/config.json"
    MAX_WAV_VALUE = 32768.0
    with open(config_file) as f:
        data = f.read()
    json_config = json.loads(data)
    h = AttrDict(json_config)

    @hk.transform_with_state
    def forward(x):
        net = Generator(h)
        return net(x)

    rng = next(hk.PRNGSequence(42))

    with open(FLAGS.ckpt_dir / "hk_hifi.pickle", "rb") as f:
        params = pickle.load(f)
    aux = {}
    wav, aux = forward.apply(params, aux, rng, mel)
    wav = jnp.squeeze(wav)
    audio = jax.device_get(wav)
    return audio