import pyDOE2
import numpy as np
import pandas as pd

from typing import Callable, Any

from IO_Objects import Descriptor
from problem import Problem



########################################################################################################
##############################     Static  sampling     ################################################
########################################################################################################


_sampler = Callable[[int, int, Any], np.ndarray]  # the format of all sampler functions


def dist_sampler(sampler: _sampler, problem: Problem, num_samples: int, *args, **kwargs) -> pd.DataFrame:
    """uses the sampling function provided to generate `num_samples` sets of values
    These values will be valid for the corresponding inputs

    :param problem: Problem that the inputs should apply to
    :param sampler: a function that is used to produce the distribution of samples
    :param num_samples: number of samples to take

    :param args: Arguments passed to the sampling function
    :param kwargs: Arguments passed to the sampling function
    :return: pandas DataFrame with one column per parameter, and `num_samples` rows

    """
    samples = sampler(num_samples, problem.num_inputs, *args, **kwargs)
    data = {name: _sample_param(param, samples[:, i]) for i, (name, param)
            in enumerate(zip(problem.names('inputs'), problem.inputs))}
    df = pd.DataFrame(data)
    # enforce the correct order in case it was lost by the dictionary
    df = df[problem.names('inputs')]
    return df


def add_extremes(df, problem):
    """adds two datapoints with the minimum and maximum values for each attribute

    :param data:
    :param problem:
    :return:
    """
    new = dist_sampler(extremes, problem, 2)
    return pd.concat([df, new], ignore_index=True)


def extremes(samples: int, attributes: int=None) -> np.ndarray:
    if attributes is None:
        samples, attributes = 2, samples
    assert samples == 2, "extremes can only produce two samples"
    return np.array([np.zeros(attributes), np.ones(attributes)])


def _sample_param(parameter: Descriptor, values):
    try:
        # this will only work if p.sample is numpy compatible
        return parameter.sample(values)
    except TypeError:
        # apply non-vectorised version of p.sample
        return [parameter.sample(x) for x in values]


# all of the following return a 2d array of shape (samples, attributes)
# which contains values between 0 and 1


random: _sampler = np.random.rand


def seeded_sampler(samples: int, attributes: int, seed=0) -> np.ndarray:
    np.random.seed(seed)
    return np.random.rand(samples, attributes)


def lhs(samples: int, attributes: int, *args, **kwargs) -> np.ndarray:
    return pyDOE2.lhs(attributes, samples, *args, **kwargs)

def full_factorial(samples: int, attributes: int, level: int=None) ->np.ndarray:
   if level == None:
       level = np.int(np.exp(np.log(samples)/attributes))
       if (level**attributes) > level:
        print(f'Total number of samples ({level**attributes}) is smaller  than input ({samples}) '
              f'to have an even number of factor levels ({level}) for all parameters.')
   return pyDOE2.fullfact((np.ones([attributes]) * level).astype(int))/level





########################################################################################################
##############################    Adaptive  sampling    ################################################
########################################################################################################


class adaptive_sampler_lv():
    """Represents an adaptive sampler which iteratively picks simulation samples and runs them."""
    def __init__(self, P, P_out, n_new, problem, evaluator, reg, scaler=None, scaler_out=None, verbose = False):
        if len(P[:,0])<2*len(P[0,:]):
            print('Size of initial number of samples needs to be bigger than 2 times no. inputs.')
        self.P = P
        self.P_out = P_out
        self.n_new = n_new
        self.problem = problem
        self.evaluator = evaluator
        self.reg = reg
        self.scaler = scaler
        self.scaler_out = scaler_out
        self.verbose = verbose

    def run(self, no_iter):

        self.score = np.empty([no_iter+1])

        self.N_P, self.S_P = init_Neighborhood(self.P)
        self.pick_new_samples()
        self.update_model()

        for i in range(no_iter):
            print(i)
            for p_new in self.P_new:
                self.N_P, self.S_P = update_Neighborhood(self.N_P, self.P, self.S_P, p_new)
                if self.verbose==True:
                    print('initialize neighborhood for new samples')
                N_P_new, S_P_new = init_Neighborhood(self.P, p_new)
                self.N_P = np.append(self.N_P, N_P_new, axis=2)
                self.S_P = np.append(self.S_P, S_P_new, axis=0)
            # update existing set of samples
            self.P = np.append(self.P, self.P_new, axis=0)
            self.P_out = np.append(self.P_out, self.P_out_new, axis=0)
            self.pick_new_samples()

            self.update_model()

            self.score[i] = self.reg.score(self.P, self.P_out)

    def update_model(self):
        self.reg.fit(self.P, self.P_out)
        if self.verbose == True:
            print(self.reg.score(self.P, self.P_out))

    def pick_new_samples(self):
        #2) Compute nonlinearity measure and Voronoi Cell
        E = LOLA_estimate(self.N_P, self.P, self.reg, scaler=self.scaler, scaler_out=self.scaler_out)
        V_P,samples = MC_Voronoi_estimate(self.P, self.problem)
        H = hybrid_score(E, V_P)
        ind_new = np.argsort(H)
        P_sorted = self.P[ind_new,:]
        if self.verbose==True:
            print('Collecting new samples around:')
            print(P_sorted[-self.n_new:,:])

        # Generate new samples in Voronoi cell taking the mean of the neighborhood
        ind=2
        self.P_new=np.empty([self.n_new,len(self.P[0,:])])
        for i in range(self.n_new):
            candidates = containedin_Voronoi(P_sorted[-i-1, :], self.P, samples)
            #print(len(candidates))
            while len(candidates)<=1:
                #print('in loop')
                candidates = containedin_Voronoi(P_sorted[-i - ind, :], self.P, samples)
                ind+=1
            self.P_new[i,:] = select_newSamp(P_sorted[-i-1, :], self.N_P[:,:,ind_new][:,:,-i-1], candidates)
        self.P_out_new = self.evaluator.df_apply(pd.DataFrame(self.P_new))




###############################################################
def hybrid_score(E,V):
    return V+E/sum(E)

def select_newSamp(p_r, N, candidates, verbose=False):
    N = np.append(N, [p_r], axis=0)
    max_dist = 0
    if verbose == True:
        print(len(candidates[1:, :]))
    for c in candidates[1:,:]:
        d=0
        for n in N:
            d = d + np.linalg.norm(c-n)
        if d>max_dist:
                max_dist = d
                max_cand = c
    return max_cand



def MC_Voronoi_estimate(P, problem):
    # P are the existing points

    # initialize arrays
    V_P = np.zeros([len(P),])  # Vector of all Voronoi sample estimates
    # per existing point collect randomly selected 100 points

    S = dist_sampler(sampling.lhs, problem, 500) # for now 100 random samples per point to estimate Voronoi cell

    for s in S.values:
        d = np.inf
        ind = 0
        for p in P:
            if np.linalg.norm(p-s) < d:
                d = np.linalg.norm(p-s)
                ind_fin = ind
            ind = ind + 1
        V_P[ind_fin] = V_P[ind_fin] + 1/len(S)
    #plt.plot(S.values[:, 0], S.values[:, 1], 'x')
    return V_P,S


def LOLA_estimate(N_P, P, model, scaler, scaler_out):
    # N_P is a 3-D array with all neighbors for each sample (mxdx#samples)
    # S_P is a 1_D array with the neighborhood score for all samples
    # P is the array including all samples

    n = len(P)
    ind = 0
    E = np.empty([n,])
    for p_r in P:
        grad = Gradient_estimation(N_P[:, :, ind], p_r, model, scaler=scaler, scaler_out=scaler_out)
        E[ind] = Nonlinearity_measure(grad, N_P[:, :, ind], p_r, model, scaler=scaler, scaler_out=scaler_out)
        ind=ind+1

    return E

def update_Neighborhood(N_P, P, S_P , P_new):
    # this function tries to add one new sample p_new to all of the existing neighborhoods
    # S is the neighborhood score
    # N_P are the associated neighbors of all points in P
    # P is the matrix of reference points
    # p_new the candidate
    if np.ndim(P)==1:
        P = np.expand_dims(P,axis=0)
        N_P = np.reshape(N_P,(len(N_P[:,0]), len(N_P[0,:]), 1))
        S_P = np.expand_dims(S_P, axis = 0)

    if np.ndim(P_new) == 1:
        P_new = np.expand_dims(P_new,axis=0)

    m = 2 * len(P[0,:])
    ind = 0
    for p_new in P_new:
        for p_r in P:
            if sum(p_r == p_new) <len(P[0,:]):
                N_temp = np.dstack([N_P[:, :, ind]]*m)
                S_temp = np.zeros([m,])
                for i in range(m):
                    # replace one sample with the new one
                    if sum(sum(N_temp[:,:,i]==p_new))<len(P[0,:]):
                        N_temp[i, :, i] = p_new
                    else:
                        pass
                    # compute neighborhood score for all possible combinations
                    S_temp[i] = Neighborhood_score(N_temp[:, :, i], p_r)
                    # if neighborhood score is better than the score without the new sample, update the neighborhood
                ind_min = np.argmin(S_temp)
                #print(ind_min)
                #print(S_temp)
                N_P[:,:,ind]=N_temp[:,:,ind_min]
                S_P[ind] = S_temp[ind_min]
                ind = ind + 1
            else:
                pass
                #print(N_P[:,:,1])
    return N_P, S_P

def init_Neighborhood(P, P_new=None, verbose=False):
    # P is the matrix with all samples
    # pnew optional if the neighborhood for new samples shall be computed
    # for now dim(pnew)=1 is assumed
    m = 2 * len(P[0,:]) # number of samples in neighborhood
    d = len(P[0,:])# number of dimensions

    if np.any(P_new==None):
        # case I: first initialization after static sampling
        n = len(P)  # number of samples
        N_P = np.empty([m,d,n])
        S_P = np.empty([n,])
        P_ref = P*1.0   # array of all samples to be considered as reference in the next for loop
    else:
        # case II: initialize neighborhoods of new samples p_new
        if np.ndim(P_new) == 1:
            if verbose==True:
                print('ndim')
            P_ref = np.expand_dims(P_new, axis=0)
        else:
            P_ref = P_new*1.0
        n_new = len(P_ref)  # number of new samples
        N_P = np.empty([m,d,n_new])
        S_P = np.empty([n_new,])


    for i in range(len(P_ref)):
        mask = P != P_ref[i,:]
        P_NoRefPoint = P[mask[:, 0], :]
        N_P[:,:,i] = P_NoRefPoint[0:m, :]
        S_P[i] = Neighborhood_score(N_P[:,:,i], P_ref[i,:])
    ind = 0
    for p_cand in P:
        ind +=1
        # at the moment p_new to in P --> only 1 p_new allowed
        N_P, S_P = update_Neighborhood(N_P, P_ref, S_P, p_cand)
    if verbose == True:
        print('Neighborhood initiated!')
    return N_P, S_P


def Neighborhood_score(N, p_r):
    # To Do: make it possible to have 3D arrays as inputs
    # N is the neighborhood
    # P_cand should have m entries
    m = len(N)

    # initialize arrays
    Dist_cand = np.empty([m,])
    minDist = np.empty([m, ])

    # compute distance matrix
    C = 0
    for i in range(m):
        C=C+np.linalg.norm(N[i,:] - p_r)
    C = C/m

    for i in range(m):
        for j in range(m):
            if i==j:
                pass
            else:
                Dist_cand[i] = np.linalg.norm(N[i,:]-N[j,:])
        minDist[i] = min(Dist_cand)

    A = 1/m*sum(minDist)
    R = A/(np.sqrt(2)*C)
    return R/C

def Gradient_estimation(N, p_r, model,scaler, scaler_out):
    m = len(N)
    d = len(p_r)

    P_mat = np.empty([m, d])
    F_mat = np.empty([m,])

    for i in range(m):
            P_mat[i, :] = N[i,:]-p_r
            if scaler == None:
                F_mat[i] = model.predict([N[i,:]])
            else:
                F_mat[i] = scaler_out.inverse_transform(model.predict(scaler.transform([N[i, :]])))
    #grad = np.linalg.solve(P_mat,F_mat.reshape((len(F_mat),1)))
    grad = np.linalg.lstsq(P_mat, np.transpose(F_mat),rcond=None)[0].reshape((1,d))
    return [grad]

def Nonlinearity_measure(grad, N, p_r, model, scaler, scaler_out):
    E=0
    for i in range(len(N)):
        if scaler == None:
            E = E + abs(model.predict([N[i, :]]) -
                        (model.predict([p_r]) + np.dot(grad,(N[i, :]-p_r))))
        else:
            E = E + abs(scaler_out.inverse_transform(model.predict(scaler.transform([N[i, :]]))) -
                        (scaler_out.inverse_transform(model.predict(scaler.transform([p_r]))) + np.dot(grad,(N[i, :]-p_r))))
    return E



def containedin_Voronoi(p_r,P,samples):
    mask = P != p_r
    P_temp = P[mask[:, 0], :]
    candidates=np.empty([1,len(P[0,:])])
    for s in samples.values:
        for p_j in P_temp:
            if np.linalg.norm(s-p_j)<=np.linalg.norm(s-p_r):
                break
            else:
                continue
        if np.all(p_j == P[-1]):
            candidates = np.append(candidates, [s], axis=0)

    return candidates
