import torch
import time
from data import load_dataset, load_multi_domain_dataset
from models import StyleTransformer, Discriminator, DomainDiscriminator
from train import train, auto_eval


class Config:
    use_domain_adaption = False
    data_path = "./data/yelp/"
    yelp_data_path = "./data/yelp/"
    imdb_data_path = "./data/imdb/"
    log_dir = "runs/exp"
    save_path = "./save"
    pretrained_embed_path = "./embedding/"
    device = torch.device(
        "cuda" if True and torch.cuda.is_available() else "cpu"
    )
    discriminator_method = "Multi"  # 'Multi' or 'Cond'
    load_pretrained_embed = False
    min_freq = 3
    max_length = 16
    embed_size = 256
    d_model = 256
    h = 4
    num_styles = 2
    num_classes = num_styles + 1 if discriminator_method == "Multi" else num_styles
    num_domains = 2
    num_domain_classes = num_domains # + 1 if discriminator_method == "Multi" else num_domains
    num_layers = 4
    batch_size = 64
    lr_F = 0.0001
    lr_D = 0.0001
    lr_domain_D = 0.0001
    L2 = 0
    iter_D = 10
    iter_domain_D = 10
    iter_F = 5
    F_pretrain_iter = 500
    log_steps = 5
    eval_steps = 25
    learned_pos_embed = True
    dropout = 0
    drop_rate_config = [(1, 0)]
    temperature_config = [(1, 0)]

    slf_factor = 0.25
    cyc_factor = 0.5
    adv_factor = 1
    domain_cns_factor = 1

    inp_shuffle_len = 0
    inp_unk_drop_fac = 0
    inp_rand_drop_fac = 0
    inp_drop_prob = 0


def main():
    config = Config()
    if config.use_domain_adaption:
        train_iters, dev_iters, test_iters, vocab = load_multi_domain_dataset(config)
    else:
        train_iters, dev_iters, test_iters, vocab = load_dataset(config)
    print("Vocab size:", len(vocab))
    model_F = StyleTransformer(config, vocab).to(config.device)
    model_D = Discriminator(config, vocab).to(config.device)
    model_domain_D = DomainDiscriminator(config, vocab).to(config.device) if config.use_domain_adaption else None
    print(config.discriminator_method)

    train(config, vocab, model_F, model_D, train_iters, dev_iters, test_iters, model_domain_D=model_domain_D)


if __name__ == "__main__":
    main()
