from Network.network import ResNet50, ResNet110, ConvNet, ConvNet_1D
from Learnware_package import *

def get_default_convnet_setting():
    net_width, net_depth, net_act, net_norm, net_pooling = 32, 3, 'relu', 'instancenorm', 'avgpooling'
    return net_width, net_depth, net_act, net_norm, net_pooling

def get_network(model, channel, num_classes, im_size=(32, 32)):
    torch.random.manual_seed(int(time.time() * 1000) % 100000)
    net_width, net_depth, net_act, net_norm, net_pooling = get_default_convnet_setting()
    if model == 'ConvNetBN':
        net = ConvNet(channel=channel,
                      num_classes=num_classes,
                      net_width=net_width,
                      net_depth=net_depth,
                      net_act=net_act,
                      net_norm='batchnorm',
                      net_pooling=net_pooling,
                      im_size=im_size)
    elif model == 'ConvNet_1D':
        net = ConvNet_1D(channel=channel,
                         num_classes=num_classes,
                         net_width=256,
                         net_depth=net_depth,
                         net_act=net_act,
                         net_norm='batchnorm',
                         net_pooling='avgpooling')
    elif model == 'ResNet50':
        net = ResNet50(channel=channel, num_classes=num_classes)
    elif model == 'ResNet110':
        net = ResNet110(channel=channel, num_classes=num_classes)

    else:
        net = None
        exit('unknown model: %s' % model)

    gpu_num = torch.cuda.device_count()
    if gpu_num > 0:
        device = 'cuda'
        if gpu_num > 1:
            net = nn.DataParallel(net)
    else:
        device = 'cpu'
    net = net.to(device)

    return net


def train(net, trainloader, optimizer, criterion, args):
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, datum in enumerate(trainloader):
        inputs = datum[0].float().to(args.device)
        targets = datum[1].long().to(args.device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        acc = 100. * correct / total

    return acc, train_loss


def test(net, testloader, criterion, args):
    net.eval()
    test_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, datum in enumerate(testloader):
            inputs = datum[0].float().to(args.device)
            targets = datum[1].long().to(args.device)
            outputs = net(inputs)
            loss = criterion(outputs, targets)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    acc = 100. * correct / total
    return acc, test_loss

def train_task_network(net, trainloader, testloader, args, opti=None):
    # net = net.to(args.device)
    '''train network parameters setting'''

    if args.model == 'PreResNet110':
        Epoch = 100
        criterion = nn.CrossEntropyLoss()
        optimizer = torch.optim.SGD(net.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)
        optimizer.load_state_dict(opti)
        net.fc = nn.Linear(256, args.num_classes).to(args.device)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Epoch)
    else:
        lr = 0.1
        Epoch = 200
        # lr_schedule = [Epoch // 2 + 1]
        criterion = nn.CrossEntropyLoss().to(args.device)
        optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Epoch)

    '''Train network'''
    best_acc = 0
    start = time.time()
    for ep in range(Epoch+1):
        train_acc, train_loss = train(net, trainloader, optimizer, criterion, args)
        test_acc, test_loss = test(net, testloader, criterion, args)
        # print('Iteration: %d, train_acc: %f, tran_loss: %f, test_acc: %f, test_loss: %f' % (ep, train_acc, train_loss, test_acc, test_loss))
        logging.info(
            f""" > Iteration: {ep}, train_acc: {train_acc}, train_loss: {train_loss}, test_acc: {test_acc}, test_loss: {test_loss}_________________\n""")
        # if ep in lr_schedule:
            # lr *= 0.1
            # optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005)
        scheduler.step()

        if test_acc > best_acc:
            best_acc = test_acc
            best_ep = ep
            best_net = net

    time_train = time.time() - start
    return best_net, time_train, train_loss, test_loss, best_acc, best_ep