import argparse
import yaml

parser = argparse.ArgumentParser(description="Mastering Style Transfer")
parser.add_argument("--model_path", type=str, default=f"models/white_box_converter.pt", help="Path to mastering converter model")
parser.add_argument("--encoder_path", type=str, default=f"models/mastering_effects_encoder.pt", help="Path to effects encoder model")
parser.add_argument("--perform_ito", action="store_true", help="Whether to perform ITO")
parser.add_argument("--optimizer", type=str, default="RAdam", help="Optimizer for ITO")
parser.add_argument("--learning_rate", type=float, default=0.001, help="Learning rate for ITO")
parser.add_argument("--num_steps", type=int, default=100, help="Number of optimization steps for ITO")
parser.add_argument("--af_weights", nargs='+', type=float, default=[0.1, 0.001, 1.0, 1.0, 0.1], help="Weights for AudioFeatureLoss")
parser.add_argument("--sample_rate", type=int, default=44100, help="Sample rate for AudioFeatureLoss")
parser.add_argument("--path_to_config", type=str, default='networks/configs.yaml', help="Path to network architecture configuration file")

args = parser.parse_args()

# load network configurations
with open(args.path_to_config, 'r') as f:
    configs = yaml.full_load(f)
args.cfg_converter = configs['TCN']['param_mapping']
args.cfg_enc = configs['Effects_Encoder']['default']