# Function to generate a clustered grid using kmedoids method

# Inputs:

#=
-simulation_in: Simulated states to be clustered

# This is the order of the simulated series
#
# Rlag   = simulation_in(:,1);                       % Interest Rate
# ylag   = simulation_in(:,2);                       % Output
# er     = simulation_in(:,3);                       % Monetary Shock
# z      = simulation_in(:,4);                       % Productivity
# g      = simulation_in(:,5);                       % Gov. expenditure

-CN: Number of clusters
-transformation: 1=Use Principal Components
=#

# Outputs:
#=
-ES_out: Simulated states
-grid_out: CLuster centers
-cindex_out: Indicator of cluster assignments
-c.counts: Number of states assigned to each cluster
=#

#--------------------------------------------------------------------------

function  cga_grid_kms(simulation_in::Array{Float64,2},CN::Int64,transformation::Int64)

obs,vars=size(simulation_in);                         # Get number of observations and variables

grid_CGA=zeros(CN,vars);                              # Preallocate matrix where Grid Points will be stored

if transformation==1;  # Preprocess data

    mu  = mean(simulation_in,dims=1);     # compute averages
    sig = std(simulation_in,dims=1);      # compute standard deviations
    MU  = kron(ones(obs,1),mu);    # Stack the vector of means (for every observation)
    SIG = kron(ones(obs,1),sig);   # Stack the vector of standard deviations (for every observation)

    X   = (simulation_in - MU)./SIG;    # This is the standarized data

    U,S,V = svd(X);              # Perform the Singular Value Decomposition of standarized data

    Z  = X*V;                      # Principal components of standarized data

    sig_z = std(Z,dims=1);                # Standard deviations of PC

    SIG_Z = kron(ones(obs,1),sig_z);

    ZZ = Z./SIG_Z;                 # Compute cluster over ZZ

    dist=pairwise(Euclidean(),ZZ')

    c = kmedoids(dist,CN,maxiter=10000);  # Cluster possitions

    SIG_Z_GRID=kron(ones(CN,1),sig_z);

    grid_CGA=ZZ[c.medoids,:].*SIG_Z_GRID

    # Revert Back the PC transformation #

    grid_CGA = grid_CGA/V;              # This incorporates the PC

    MU_GRID  = kron(ones(CN,1),mu);    # Stack the vector of means (for every grid point)
    SIG_GRID = kron(ones(CN,1),sig);   # Stack the vector of standard deviations


    grid_CGA   = grid_CGA.*SIG_GRID + MU_GRID;

    ES_out  = simulation_in;
    grid_out = grid_CGA;                # Cluster Grid points (Over ZZ)


elseif transformation==0

    # No transformation clusters the data directly over the untransformed
    # data

    X = simulation_in;                                  # This is the ergodic set

    dist=pairwise(Euclidean(),X')

    c = kmedoids(dist,CN,maxiter=10000);  # Cluster possitions

    grid_CGA=X[c.medoids,:]

    grid_out   = grid_CGA;                   # Cluster Grid points

end


cindex_out = c.assignments;                          # Cluster index for color plotting

return grid_out, cindex_out, c.counts
end
