import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader


class SynDataset(Dataset):
    def __init__(self, X, y):
        self.X = [torch.tensor(item, dtype=torch.float32) for item in X]
        self.y = [torch.tensor(item, dtype=torch.float32, requires_grad=False) for item in y]
        
    def __len__(self):
        return len(self.X)
    
    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]
    

def get_synthetic_dataloader(args, return_idx):
    in_dim = args.in_dim
    out_dim = 1 # TODO
    sigma_y = 0.1
    num_samples_tot = 10000 + 2000 # train+test
    num_samples = num_samples_tot // args.n_workers
    n_train = 10000 // args.n_workers
    n_test = num_samples - n_train
    
    train_x_split, train_y_split = [], []
    test_x_split, test_y_split = [], []

    if args.partition == 'homo':
        # model parameters
        w = np.random.normal(0, 1, in_dim)

        # data inputs
        X = np.random.normal(0, 1, (num_samples_tot, in_dim, ))

        # data outputs
        y = np.dot(X, w) + np.random.normal(0, sigma_y, num_samples_tot)
        y = y.reshape(-1, 1)
        
        X = X.tolist()
        y = y.tolist()

        # train and test
        for i in range(args.n_workers):
            idx_begin, idx_end = i*num_samples, (i+1)*num_samples
            train_x_split.append(X[idx_begin : idx_begin+n_train])
            train_y_split.append(y[idx_begin : idx_begin+n_train])
            
            test_x_split.append(X[idx_begin+n_train : idx_end])
            test_y_split.append(y[idx_begin+n_train : idx_end])
    else:
        alpha = 1.0/args.alpha # NOTE
        mean_w = np.random.normal(0, alpha, args.n_workers) # control heterogeneity

        # model parameters
        w_list = []
        for i in range(args.n_workers):
            w = np.random.normal(mean_w[i], 1, in_dim)
            w_list.append(w.tolist())
        
        # data inputs and outputs
        for i in range(args.n_workers):
            X = np.random.normal(0, 1, (num_samples, in_dim))
            
            y = np.dot(X, w_list[i]) + np.random.normal(0, sigma_y, num_samples)
            y = y.reshape(-1, 1)

            X = X.tolist()
            y = y.tolist()
            
            train_x_split.append(X[:n_train])
            train_y_split.append(y[:n_train])
            
            test_x_split.append(X[n_train:])
            test_y_split.append(y[n_train:])
    
    if return_idx == -1:
        train_x_split = np.array(train_x_split).reshape(-1, in_dim)
        train_y_split = np.array(train_y_split).reshape(-1, 1)
        train_ds = SynDataset(train_x_split, train_y_split)
        train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True)
        
        test_x_split = np.array(test_x_split).reshape(-1, in_dim)
        test_y_split = np.array(test_y_split).reshape(-1, 1)
        test_ds = SynDataset(test_x_split, test_y_split)
        test_dl = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False)
        return train_dl, test_dl
    
    else:
        train_ds = SynDataset(train_x_split[return_idx], train_y_split[return_idx])
        train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True)
        return train_dl

# if __name__ == '__main__':
#     args =  ARGS()

#     X_split, y_split = get_synthetic_dataloader(args)
#     print(X_split[0][:5])
#     print(y_split[0][:5])

#     model = QuadraticModel(50)
#     print(model.W.shape)

#     for i in range(5):
#         x = torch.tensor(X_split[i], dtype=torch.float32)
#         y = torch.tensor(y_split[i], dtype=torch.float32).view(-1, 1)
#         output = model(x)

#         loss = nn.MSELoss()(output, y)
#         print(loss)