import os
import math
import sys
import json
import torch
import torch.nn as nn
import librosa
import numpy as np

from . import spec_utils, nets, nets_new
from .model_param_init import ModelParameters

VOCAL_STEM = "vocals"
INST_STEM = "instrumental"
OTHER_STEM = "other"
BASS_STEM = "bass"
DRUM_STEM = "drums"
GUITAR_STEM = "guitar"
PIANO_STEM = "piano"
SYNTH_STEM = "synthesizer"
STRINGS_STEM = "strings"
WOODWINDS_STEM = "woodwinds"
BRASS_STEM = "brass"
WIND_INST_STEM = "wind_inst"

NON_ACCOM_STEMS = (
    VOCAL_STEM,
    OTHER_STEM,
    BASS_STEM,
    DRUM_STEM,
    GUITAR_STEM,
    PIANO_STEM,
    SYNTH_STEM,
    STRINGS_STEM,
    WOODWINDS_STEM,
    BRASS_STEM,
    WIND_INST_STEM,
)

def get_model(config):
    model_params = ModelParameters(dict(config.model.model_params))
    nn_arch_sizes = [
        31191,
        33966,
        56817,
        123821,
        123812,
        129605,
        218409,
        537238,
        537227,
    ]
    vr_5_1_models = [56817, 218409]
    is_vr_51_model = False
    model_capacity = 32, 128
    if config.model.nout and config.model.nout_lstm:
        model_capacity = config.model.nout, config.model.nout_lstm
        is_vr_51_model = True
    if config.model.nn_arch_size in vr_5_1_models or is_vr_51_model:
        model = nets_new.CascadedNet(
            model_params.param["bins"] * 2,
            config.model.nn_arch_size,
            nout=model_capacity[0],
            nout_lstm=model_capacity[1],
        )
    else:
        model = nets.determine_model_capacity(
            model_params.param["bins"] * 2, config.model.nn_arch_size
        )
    model.model_params = model_params
    return model