import os
import random
from typing import Optional
from argparse import Namespace

import timm
import torch
import torch.nn as nn
import torch.multiprocessing as torch_mp
from torch.utils.data import DataLoader
import nni
import logging
import numpy as np

from model.transcrowd_gap import VisionTransformerGAP
from arguments import args, ret_args

logger = logging.getLogger("train")


def setup_process_group(
        rank: int,
        world_size: int,
        master_addr: str = "localhost",
        master_port: Optional[np.ushort] = None
):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = (
        str(random.randint(40000, 65545))
        if master_port is None
        else str(master_port)
    )

    # join point!
    torch.distributed.init_process_group(
        backend="nccl", rank=rank, world_size=world_size
    )

# TODO:
# The shape for each batch in transcrowd is [3, 384, 384],
# this is due to images being cropped before training.
# To preserve image semantics wrt the entire layout, we want to apply cropping
# i.e., as encoder input during the inference/training pipeline.
# This should be okay since our transformations are all deterministic?
# not sure...




def build_train_loader():
    pass


def build_valid_loader():
    pass


def train_one_epoch(
        train_loader: DataLoader,
        model: VisionTransformerGAP,
        criterion,
        optimizer,
        scheduler,
        epoch: int,
        args: Namespace
):
    # Get learning rate
    curr_lr = optimizer.param_groups[0]["lr"]
    print("Epoch %d, processed %d samples, lr %.10f" %
        (epoch, epoch * len(train_loader.dataset), curr_lr)
    )

    # Set to train mode (perspective estimator only)
    revpers_net.train()
    end = time.time()

    # In one epoch, for each training sample
    for i, (fname, img, gt_count) in enumerate(train_loader):
        # move stuff to device
        # fpass (revpers)
        img = img.cuda()

        # loss wrt revpers
        loss = criterion()


    pass

def valid_one_epoch():
    pass

def main(rank: int, args: Namespace):
    pass

if __name__ == "__main__":
    tuner_params = nni.get_next_parameter()
    logger.debug("Generated hyperparameters: {}", tuner_params)
    combined_params = Namespace(
        nni.utils.merge_parameter(ret_args, tuner_params)
    ) # Namespaces have better ergonomics, notably a struct-like access syntax.
    logger.debug("Parameters: {}", combined_params)

    if combined_params.use_ddp:
        # Use DDP, spawn threads
        torch_mp.spawn(
            main,
            args=(combined_params, ), # rank supplied automatically as 1st param
            nprocs=combined_params.world_size,
        )
    else:
        # No DDP, run in current thread
        main(None, combined_params)