import numpy as np

def linear_regression(x,y,xnew=[],sort=False):
 
  from scipy import interpolate,polyfit,polyval

  ####---- Performing a simple linear regression of y points ----####
  
  # Calculating Regression coefficient
  (ar,br)=polyfit(x,y,1);
  y_fit = polyval([ar,br],x);
  R = np.corrcoef(y_fit,y)[0,1]

  # Calculating regression line
  if ( len(xnew)<>0 ):
    y_fit = polyval([ar,br],xnew);
  else:
    xnew = x;
    y_fit = polyval([ar,br],xnew);

  # Sorting array
  if sort == True :
    idx_sort = np.argsort(xnew);
    xnew = xnew[idx_sort]; y_fit = y_fit[idx_sort];

  return y_fit,xnew,R


def reading_file_topo2d(filename='topo2d',verb=True):
  
  ####---- Reads topo2d-type output file and assemble into a t-x grid ----####
  
  if (verb==True): print("... Reading %s "%filename)
  f = open(filename,'r') 

  x_grid = [];
  z_grid = [];
  t_grid = [];
  vz_grid = [];

  a = [];
  for row in f: 
    a.append(row); 
  f.close(); 
  nouts=-1
  for i in a: 
    x=i.split();
    if (x[0]=="#"):
     nouts+=1; time = x[1];
     x_grid.append([]); z_grid.append([]); t_grid.append([]); vz_grid.append([]);
    else:
      x_grid[nouts].append(float(x[0])); z_grid[nouts].append(float(x[1]));
      t_grid[nouts].append(float(time)); vz_grid[nouts].append(float(x[4]))

  ntime = np.shape(z_grid)[0]
  nptsx = np.shape(z_grid)[1]

  if (verb == True): print(" \t nptsx = %d ; ntimes = %d"%(nptsx,ntime))
  if (verb == True): print(" \t Min/Max times : %5.2d/%5.2f"%(np.min(t_grid),np.max(t_grid)))

  return np.array(t_grid),np.array(x_grid),np.array(z_grid),np.array(vz_grid)


def indexes_times_to_extract(times_to_extract,t_array):
  ##-- Outputs the elements indexes corresponding to the times to extract in t_array
  idx_to_extract = [];
  for t in times_to_extract:
    idx_to_extract.append(np.argmin(np.abs(t-t_array)))
  return idx_to_extract



def reading_file_xy(filename='xy.tab',verb=True):

  ####---- Reads a simple two-column file and place into numpy-arrays x and y

  if (verb == True): print("... In reading_file_xy \n\t... Reading %s "%filename)
  f = open(filename,'r')
  
  x_array=[]; y_array=[];
  
  a = [];
  for row in f:
    x = row.split()[0]; y = row.split()[1];
    x_array.append(float(x)); y_array.append(float(y));
  f.close();

  return np.array(x_array),np.array(y_array)


def smoothing_1d_samex(X,Y_old,c_smooth,kspline):

  from scipy import interpolate

  if (c_smooth <> None):
    var=np.var(Y_old);
    smooth=len(Y_old)*var*c_smooth

  spline=interpolate.UnivariateSpline(X,Y_old,k=kspline,s=smooth)
  Y_new=spline(X)
  

  return Y_new


def smoothing_grid_splines(grid,axis_to_smooth,kspline=3,coeffs=0.0,verb=True):
  
  ####---- Smoothing a 2-d grid along one axis using UnivariateSpline
  
  axis_to_loop = 0 if axis_to_smooth == 1 else 1 if axis_to_smooth == 0 else nan
  if (verb==True): print("\n \t Smoothing along axis = %1d"%axis_to_loop);
  if (verb==True): print("\t...  thus slicing %2d times"%np.shape(grid)[axis_to_loop])
  
  if type(grid).__module__ != 'numpy':
    sys.exit(" \n in %s: grid is not numpy.array" \
        %(sys._getframe().f_code.co_name) )

  grid_smoothed = np.zeros(np.shape(grid));
  for k in range(0,np.shape(grid)[axis_to_loop]):
    y_data = np.take(grid,indices=k,axis=axis_to_loop);
    x = np.arange(0,np.shape(grid)[axis_to_smooth]);
    y_smooth = smoothing_1d_samex(x,y_data,coeffs,kspline);
    if axis_to_loop == 0:
      grid_smoothed[k,:]=y_smooth;
    elif axis_to_loop == 1:
      grid_smoothed[:,k]=y_smooth;

  return grid_smoothed


def minmax_DH_DHr_after_tvar(t_grid,z_grid,z_grid_rate,idx_tvar,idx_point,tmax):

  ####---- Uses grid files to calculate the min/max of points with index idx_point in the grid ----####

  itmax = np.argmin(np.abs(t_grid[:,0]-tmax)); ## Index of tmax in tarray
  
  #-- Min of Elevation
  DHmin  = np.min(z_grid[np.int(idx_tvar):itmax+1,idx_point]-z_grid[idx_tvar,idx_point]);
  itDHmin = np.argmin(z_grid[np.int(idx_tvar):itmax+1,idx_point]-z_grid[idx_tvar,idx_point]); tDHmin = t_grid[np.int(idx_tvar):itmax+1,0][itDHmin];
  #-- Max of Elevation
  DHmax  = np.max(z_grid[np.int(idx_tvar):itmax+1,idx_point]-z_grid[idx_tvar,idx_point]);
  itDHmax = np.argmax(z_grid[np.int(idx_tvar):itmax+1,idx_point]-z_grid[idx_tvar,idx_point]); tDHmax = t_grid[np.int(idx_tvar):itmax+1,0][itDHmax];

  #-- Min of Rate Elevation
  DHrmin = np.min(z_grid_rate[np.int(idx_tvar)-1:itmax+1,idx_point]);
  itDHrmin = np.argmin(z_grid_rate[np.int(idx_tvar):itmax+1,idx_point]);
  if (itDHrmin<len(t_grid[np.int(idx_tvar):itmax+1,0])-1): tDHrmin = t_grid[np.int(idx_tvar):itmax+1,0][itDHrmin+1];
  else: tDHrmin = t_grid[np.int(idx_tvar):itmax+1,0][itDHrmin];
  #-- Max of Rate Elevation
  DHrmax = np.max(z_grid_rate[np.int(idx_tvar)-1:itmax+1,idx_point]);
  itDHrmax = np.argmax(z_grid_rate[np.int(idx_tvar):itmax+1,idx_point]);
  if (itDHrmax<len(t_grid[np.int(idx_tvar):itmax+1,0])-1): tDHrmax = t_grid[np.int(idx_tvar):itmax+1,0][itDHrmax+1];
  else: tDHrmax = t_grid[np.int(idx_tvar):itmax+1,0][itDHrmax];

  return DHmin,itDHmin,tDHmin, DHmax, itDHmax, tDHmax, DHrmin, itDHrmin, tDHrmin, DHrmax, itDHrmax, tDHrmax



def find_bulges_depressions(xarr,zarr,xmax_for_d1,xmax_for_d2):
  
  ####---- Find top and bottom of bulges and depression of a topo profile  ----####
  
  idx_to_start = 1; # to find local min/max other than trench

  idx_xmax1     = np.argmin(np.abs(xarr-xmax_for_d1)); # index of xmax
  idx_xmax2     = np.argmin(np.abs(xarr-xmax_for_d2)); # index of xmax
  
  ###--- 1st depression is lowest minimum after trench  <--- NEED TO BE DETERMINED FIRST
  #-- Find all local minima
  a = np.r_[True, zarr[1:idx_xmax1] < zarr[:idx_xmax1-1]] & np.r_[zarr[:idx_xmax1-1] < zarr[1:idx_xmax1], True]; idx_loc_min = np.where(a==True);
  idx_loc_min_d1 = np.argmin([x for i,x in enumerate(zarr[:idx_xmax1][idx_loc_min[0][idx_to_start:]])  ]);
  idx_d1 = idx_loc_min[0][idx_to_start:][idx_loc_min_d1]; idx_d1_glob = np.argmin( np.abs(xarr - xarr[:idx_xmax1][idx_d1])  );
  xd1 = xarr[idx_d1_glob]; zd1 = zarr[idx_d1_glob];

  ###--- 1st bulge is highest local maximum between trench and 1st depression
  #-- Find all local maxima between trench and 1st depression
  a = np.r_[True, zarr[1:idx_d1] > zarr[:idx_d1-1]] & np.r_[zarr[:idx_d1-1] > zarr[1:idx_d1], True]; idx_loc_max = np.where(a==True);
  if (len(idx_loc_max[0])>1):
    idx_loc_max_b1 = np.argmax([x for i,x in enumerate(zarr[:idx_d1][idx_loc_max[0][:]])  ]);
    idx_b1 = idx_loc_max[0][:][idx_loc_max_b1];
  else:
    idx_loc_max_b1 = 0;
    idx_b1 = idx_loc_max[0];
  xb1 = xarr[:idx_d1][idx_b1]; zb1 = zarr[:idx_d1][idx_b1];

  ####--- 2nd bulge is highest local maximum between 1st depression and last point
  #-- Find all local maxima between 1st depression and last point
  #--   --> Constraint : 2nd bulge has to be before imposed xmax for 2nd depre
  a = np.r_[True, zarr[idx_d1+1:idx_xmax2] > zarr[idx_d1:idx_xmax2-1]] & np.r_[zarr[idx_d1:idx_xmax2-1] > zarr[idx_d1+1:idx_xmax2], True]; idx_loc_max = np.where(a==True);
  idx_loc_max_b2 = np.argmax([x for i,x in enumerate(zarr[idx_d1:idx_xmax2][idx_loc_max[0][:]])  ]);
  idx_b2 = idx_loc_max[0][:][idx_loc_max_b2]; idx_b2_glob = np.argmin( np.abs(xarr - xarr[idx_d1:idx_xmax2][idx_b2])  );
  xb2 = xarr[idx_d1:][idx_b2]; zb2 = zarr[idx_d1:][idx_b2];


  ###--- Forearc point defined as middle point between trench and 1st bulge
  xtrench = xarr[0]
  xoforarc = (xb1 + xtrench)/2.0;
  xiforarc = (xd1 + xb1)/2.0;
  xmidarc  = (xb2 + xd1)/2.0;
  idx_oforarc_glob = np.argmin( np.abs(xarr - xoforarc) );
  idx_iforarc_glob = np.argmin( np.abs(xarr - xiforarc) );
  idx_midarc_glob = np.argmin( np.abs(xarr - xmidarc) );
  zoforarc = zarr[idx_oforarc_glob];
  ziforarc = zarr[idx_iforarc_glob];
  zmidarc  = zarr[idx_midarc_glob];

  return xoforarc,zoforarc,xiforarc,ziforarc,xmidarc,zmidarc