import os
import torch
import numpy as np
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.nn.functional as F
import matplotlib.pyplot as plt
def pca_show(feats):
    feats = feats.cpu().detach()
    B, C, H, W = feats.shape
    n_components = 3
    from sklearn.decomposition import PCA
    pca = PCA(n_components=n_components)      
    feats = F.interpolate(feats, size=(128, 128), mode='bilinear', align_corners=False)  
    features = feats[0, :, :, :].squeeze().reshape(C,-1).permute(1,0).cpu()
    pca.fit(features)
    pca_features = pca.transform(features)
    pca_features = (pca_features - pca_features.min()) / (pca_features.max() - pca_features.min())
    pca_features = pca_features * 255
    pca_features = pca_features.reshape(128, 128, n_components).astype(np.uint8)
    plt.imshow(pca_features)


def loadmodel(model, filename, strict=True, remove_prefix=False):
    if os.path.exists(filename):
        params = torch.load('%s' % filename)
        if remove_prefix:
            new_params = {k.replace('module.', ''): v for k, v in params.items()}
            model.load_state_dict(new_params,strict=strict)
            print("prefix successfully removed")
        else:
            model.load_state_dict(params,strict=strict)
        print('Load %s' % filename)
    else:
        print('Model Not Found')
    return model

def loadoptimizer(optimizer, filename):
    if os.path.exists(filename):
        params = torch.load('%s' % filename)
        optimizer.load_state_dict(params)
        print('Load %s' % filename)
    return optimizer

def loadscheduler(scheduler, filename):
    if os.path.exists(filename):
        params = torch.load('%s' % filename)
        scheduler.load_state_dict(params)
        print('Load %s' % filename)
    else:
        print('Scheduler Not Found')
    return scheduler

def savemodel(model, filename):
    print('Save %s' % filename)
    torch.save(model.state_dict(), filename)

def saveoptimizer(optimizer, filename):
    print('Save %s' % filename)
    torch.save(optimizer.state_dict(), filename)

def savescheduler(scheduler, filename):
    print('Save %s' % filename)
    torch.save(scheduler.state_dict(), filename)

def optimizer_setup_Adam(net, lr = 0.0001, init=True, step_size=3, stype='step'):
    print(f'Optimizer (Adam) lr={lr}')
    if init==True:
        net.init_weights()
    net = torch.nn.DataParallel(net)
    optim_params = [{'params': net.parameters(), 'lr': lr},] # confirmed
    optimizer = torch.optim.Adam(optim_params, betas=(0.9, 0.999), weight_decay=0)
    if stype == 'cos':
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 30, eta_min=0, last_epoch=-1)
        print('Cosine aneealing learning late scheduler')
    if stype == 'step':
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.8)
        print(f'Step late scheduler x0.8 decay every {step_size}')
    return net, optimizer, scheduler

def optimizer_setup_SGD(net, lr = 0.01, momentum= 0.9, init=True):
    print(f'Optimizer (SGD with momentum) lr={lr}')
    if init==True:
        net.init_weights()
    net = torch.nn.DataParallel(net)
    optim_params = [{'params': net.parameters(), 'lr': lr},] # confirmed
    return net, torch.optim.SGD(optim_params, momentum=momentum, weight_decay=1e-4, nesterov=True)

def optimizer_setup_AdamW(net, lr = 0.001, eps=1.0e-8, step_size = 20, init=True, stype='step', use_data_parallel=False):
    print(f'Optimizer (AdamW) lr={lr}')
    if init==True:
        net.init_weights()
    if use_data_parallel:
        net = torch.nn.DataParallel(net)
    optim_params = [{'params': net.parameters(), 'lr': lr},] # confirmed
    optimizer = torch.optim.AdamW(optim_params, betas=(0.9, 0.999), eps=eps, weight_decay=0.01)
        
    if stype=='cos':
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 30, eta_min=0, last_epoch=-1)
        print('Cosine aneealing learning late scheduler')
    if stype == 'step':
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.8)
        print(f'Step late scheduler x0.8 decay every {step_size}')
    return net, optimizer, scheduler

def mode_change(net, Training):
    if Training == True:
        for param in net.parameters():
            param.requires_grad = True
        net.train()
    if Training == False:
        for param in net.parameters():
            param.requires_grad = False
        net.eval()


def get_n_params(model):
    pp=0
    for p in list(model.parameters()):
        nn=1
        for s in list(p.size()):
            nn = nn*s
        pp += nn
    return pp

def loadCheckpoint(path, model, cuda=True):
    if cuda:
        checkpoint = torch.load(path)
    else:
        checkpoint = torch.load(path, map_location=lambda storage, loc: storage)
    model.load_state_dict(checkpoint['state_dict'])

def saveCheckpoint(save_path, epoch=-1, model=None, optimizer=None, records=None, args=None):
    state   = {'state_dict': model.state_dict(), 'model': args.model}
    records = {'epoch': epoch, 'optimizer':optimizer.state_dict(), 'records': records,
            'args': args}
    torch.save(state, os.path.join(save_path, 'checkp_%d.pth.tar' % (epoch)))
    torch.save(records, os.path.join(save_path, 'checkp_%d_rec.pth.tar' % (epoch)))



def masking(img, mask):
    # img [B, C, H, W]
    # mask [B, 1, H, W] [0,1]
    img_masked = img * mask.expand((-1, img.shape[1], -1, -1))
    return img_masked

def print_model_parameters(model):
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])


    print('# parameters: %d' % params)


def angular_error(x1, x2, mask = None): # tensor [B, 3, H, W]

    if mask is not None:
        dot = torch.sum(x1 * x2 * mask, dim=1, keepdim=True)
        dot = torch.max(torch.min(dot, torch.Tensor([1.0-1.0e-12])), torch.Tensor([-1.0+1.0e-12]))
        emap = torch.abs(180 * torch.acos(dot)/np.pi) * mask
        mae = torch.sum(emap) / torch.sum(mask)
        return mae, emap
    if mask is None:
        dot = torch.sum(x1 * x2, dim=1, keepdim=True)
        dot = torch.max(torch.min(dot, torch.Tensor([1.0-1.0e-12])), torch.Tensor([-1.0+1.0e-12]))
        error = torch.abs(180 * torch.acos(dot)/np.pi)
        return error

def write_errors(filepath, error, trainid, numimg, objname = []):
    dt_now = datetime.datetime.now()
    print(filepath)

    if len(objname) > 0:
        with open(filepath, 'a') as f:
            f.write('%s %03d %s %02d %.2f\n' % (dt_now, numimg, objname, trainid, error))
    else:
        with open(filepath, 'a') as f:
            f.write('%s %03d %02d %.2f\n' % (dt_now, numimg, trainid, error))


def save_nparray_as_hdf5(self, a, filename):
    h5f = h5py.File(filename, 'w')
    h5f.create_dataset('dataset_1', data=a)
    h5f.close()

def freeze_params(net):
    for param in net.parameters():
        param.requires_grad = False

def unfreeze_params(net):
    for param in net.parameters():
        param.requires_grad = True

def make_index_list(maxNumImages, numImageList):
    index = np.zeros((len(numImageList) * maxNumImages), np.int32)
    for k in range(len(numImageList)):
        index[maxNumImages*k:maxNumImages*k+numImageList[k]] = 1
    return index