import numpy as np
import torch
import random
import pickle
import os
from tools import optimization_oracle_tools
from tools import iceo_framework
from sklearn.model_selection import train_test_split

def random_simplex(dim_cost, n):
    alpha = np.ones(dim_cost)
    # Sample from Dirichlet distribution
    samples = np.random.dirichlet(alpha, n)
    return samples

def pw_dataset(problem_params):
    n = problem_params['oracle_sample_size']
    K = problem_params['dim_scenario']
    p_train = random_simplex(K, n)
    print(p_train)
    w_train = [optimization_oracle_tools.opt_oracle(prob_vec, problem_params) for prob_vec in p_train]
    p_train = np.array(p_train)
    w_train = np.array(w_train)
    p_val = random_simplex(K, 500)
    w_val = [optimization_oracle_tools.opt_oracle(prob_vec, problem_params) for prob_vec in p_val]
    p_val = np.array(p_val)
    w_val = np.array(w_val)
    # print(p_train)
    p_test = random_simplex(K, 500)
    w_test = [optimization_oracle_tools.opt_oracle(prob_vec, problem_params) for prob_vec in p_val]
    p_test = np.array(p_val)
    w_test = np.array(w_val)
    # print("p_samples array shape: ", np.shape(p_array), "\n")
    # print("w_samples array shape: ", np.shape(w_array), "\n")
    # p_train, p_test, w_train, w_test = train_test_split(p_array, w_array, test_size = 0.2, random_state = 66)
    return p_train, p_val, p_test, w_train, w_val, w_test

def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x /np.sum(e_x)


def generate_true_hypothesis(problem_params, seed=None):
    dim_features = problem_params['dim_features']
    dim_scenario = problem_params['dim_scenario']
    if seed is not None:
        np.random.seed(seed)
    # b_out = 0.04 * np.random.rand(dim_scenario, dim_features) - 0.02
    b_out = 0.04 * np.random.rand(dim_scenario, dim_features) - 0.02
    if problem_params['data_generation_model'] == 'context_shift':
        np.random.seed(6)
        b_out =150*np.random.rand(dim_scenario, dim_features)
    return b_out

def generate_true_hypothesis_NN(problem_params, seed=123):
    torch.manual_seed(seed)
    dim_features = problem_params['dim_features']
    dim_scenario = problem_params['dim_scenario']
    if problem_params['data_generation_model'] == 'multi_layer':
        model = iceo_framework.multi_layers_true(dim_features, dim_scenario, 128)
    else:
        model = iceo_framework.two_layers_true(dim_features, dim_scenario, 128)
    return model



def xz_dataset(size, problem_params, B_true, shift=False):
    M = problem_params['x_range']
    dim_feature = problem_params['dim_features']
    dim_scenario = problem_params['dim_scenario']
    dim_sol = problem_params['dim_sol']
    scena_list = problem_params['scenario_list']
    deg = problem_params['deg_data']
    if problem_params['data_generation_model'] == 'model mis-specification-deg':
        x_samples = M * np.random.normal(loc=0, scale=M, size=(size, dim_feature))
        z_samples = []
        for i in range(size):
            # x_samples.append( np.concatenate((x_t[i], np.array([1]))))
            p_true = softmax((np.matmul(B_true, x_samples[i]))**deg)
            #print("p_true", p_true)
            index = np.random.choice(dim_scenario, 1, p=p_true.tolist())
            # print("index", index)
            z_tilde = scena_list[index]
            z_samples.append(z_tilde)
        z_samples = np.reshape(z_samples, (size, dim_sol))
        x_samples = np.reshape(x_samples, (size, dim_feature))
        return x_samples, z_samples
    else: 
        x_samples = M * np.random.normal(loc=0, scale=M, size=(size, dim_feature))
        z_samples = []
        x_samples_tensor = torch.as_tensor(x_samples, dtype=torch.float32)
        p_true_tensor = B_true(x_samples_tensor)
        p_true = p_true_tensor.detach().numpy()
        for i in range(size):
           p_true_vec = p_true[i].astype('float64')
           p_true_vec /= p_true_vec.sum()
        #    print(np.sum(p_true_vec))
           index = np.random.choice(dim_scenario, 1, p = p_true_vec.tolist())
           #print("index", index)
           z_tilde = scena_list[index]
           z_samples.append(z_tilde)
        z_samples = np.reshape(z_samples, (size, dim_sol))
        x_samples = np.reshape(x_samples, (size, dim_feature))
        #print(z_samples)
        return x_samples, z_samples
    




