import numpy as np

import scipy

import nibabel as nib

import matplotlib.pyplot as plt

from scipy.spatial import Delaunay

import matplotlib

matplotlib.use('Agg')

from scipy.interpolate import RegularGridInterpolator

# this script creates a mesh with triangulation density locally-adaptive based on an average surface area.

# the original grid spacing was 128 points over 20mm (or 20/128)

# since surf vertex area is proportional to grid spacing squared,

# if S_target and S_in are the target and input surface areas, then,

# adjustment to get corrected grid spacing: 

# (20/N)^2 = (S_target / S_in)  * (20/128)^2

# equiv to:

# N = 128 * sqrt(S_in / S_target)

# the approach this script takes is to create a set of grids at diff resolutions, and pick the grid points based on the binned input surface area..

# all the points from the different resolutions are then triangulated

target_surfarea = snakemake.params.targetarea

num_bins = snakemake.params.nbins

n_ap = snakemake.config['in_surfarea']['dims'][0]

n_pd = snakemake.config['in_surfarea']['dims'][1]

start_ap = snakemake.config['in_surfarea']['start'][0]

end_ap = snakemake.config['in_surfarea']['end'][0]

start_pd = snakemake.config['in_surfarea']['start'][1]

end_pd = snakemake.config['in_surfarea']['end'][1]

N_pd = n_pd+2 # i.e. 128 for hipp, 32 for dentate

aspect_ratio = int((n_ap+2) / (n_pd+2)) # e.g. 256/128 = 2 for hipp, and 256/32 = 5 for dentate

#load average surface area metric

surfarea_gii = nib.load(snakemake.input.surfarea_gii)

arr_surfarea = surfarea_gii.get_arrays_from_intent('NIFTI_INTENT_NORMAL')[0].data.reshape(n_pd,n_ap)

#replace nan with 0.01 (actually, shouldn't be any nans with the latest workflow) 

arr_surfarea = np.nan_to_num(arr_surfarea,nan=0.01)

#bin the image into discrete regions using histogram

(histval,histedges) = np.histogram(arr_surfarea.flat,bins=num_bins)

binned_area = np.zeros(arr_surfarea.shape)

for i in range(len(histedges)-1):

    masked = np.logical_and(arr_surfarea>=histedges[i],arr_surfarea<histedges[i+1])

    binned_area[masked] = i

print(binned_area.shape)

# create an interpolator to sample the surfarea on each new grid

orig_x = np.linspace(start_ap,end_ap,n_ap) #since the original flat space was offset a bit from 0-40,0-20

orig_y = np.linspace(start_pd,end_pd,n_pd)

print(orig_x)

print(orig_y)

#now, for each multi-res grid, interpolate the bin value

interpolator = RegularGridInterpolator((orig_y, orig_x), binned_area,method='nearest')

multires_points = list()

#create grid for each bin centre

for i in range(len(histedges)-1):

    in_surfarea = (histedges[i] + histedges[i+1]) * 0.5 # histogram bin centre 

    #bin_centre is the mean vertex area for those vertices

    # adjustment to get corrected grid spacing (1/N)

    # 1/N = (S_target / S_in)^2  * (1/128)

    # equiv to:

    # N = 128 * sqrt(S_in / S_target)

    print(f'input surf area at bin {i}: {in_surfarea}')

    N = int(N_pd * np.sqrt(in_surfarea/ target_surfarea))

    #N = int(snakemake.params.scaling_factor * np.power(in_surfarea / target_surfarea,snakemake.params.power_factor))

    print(f'N for gridding is: {N}')

    print(N)

    nx, ny = (aspect_ratio*N,N)

    x = np.linspace(start_ap,end_ap, nx)

    y = np.linspace(start_pd,end_pd, ny)

    #print(x)

    #print(y)

    xv, yv = np.meshgrid(y,x)

   # print(xv.max())

   # print(yv.max())

    binval = interpolator((xv,yv))

    #now get points for the binval

    #coords = np.unravel_index(np.where(binval==i),xv.shape)

    #print(coords)

    coords_x = xv[np.where(binval==i)]

   # print(coords_x)

    coords_y = yv[np.where(binval==i)]

   # print(coords_y)

    multires_points.append(np.transpose(np.vstack((coords_x,coords_y))))

all_points = np.vstack(multires_points)

print(all_points.shape)

#plot a grid

tri = Delaunay(all_points)

plt.figure(figsize=(50,100))

plt.triplot(all_points[:,0], all_points[:,1], tri.simplices)

plt.plot(all_points[:,0], all_points[:,1], '.')

plt.savefig(snakemake.output.grid_png)

#write the vertices and triangles as csv 

#so we can create gifti from matlab

nverts = all_points.shape[0]

print(f'nverts: {nverts}')

nbins = num_bins

vertices_fname = snakemake.output.points_csv

triangles_fname = snakemake.output.triangles_csv

np.savetxt(vertices_fname,tri.points,fmt='%f')

np.savetxt(triangles_fname,tri.simplices,fmt='%u')