#import necessary packages
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import os

from itertools import product
from matplotlib.ticker import NullFormatter
from numpy.linalg import norm
from scipy.optimize import nnls


def em_mixing_model(samdf, emdf, Sig_species, n_iter = 10000, rmse_cutoff = 0.01, 
	pyEMs = ['P'], pptdf = None):
	'''
	Performs end-member mixing model.

	Parameters
	----------
	samdf : pd.DataFrame
		DataFrame of measured sample data. Must contain columns:
			
			Ca_uM, Cl_uM, K_uM, Mg_uM, Na_uM, SO4_uM,

		which are used to perform mixing model. Index names correspond to 
		sample names.

	emdf : pd.DataFrame
		DataFrame of inputted end-member data. Must contain rows:

			ChiCa, ChiCl, ChiMg, ChiNa, ChiS, ChiK,

		which are the fractional charge contribution of each species (i.e., the
		charge divided by the sum of cation charge or dividied by the sum of
		cation + sulfate charge, if pyrite is to be a separate end-member). 
		Must contain columns:

			X_min, X_max, etc.

		where "X" is a letter corresponding to each end-member of interest (e.g.,
		"L_min", "L_max" for limestone, etc.)

	Sig_species : list
		List of species that are included in the denominator of all scalings,
		e.g., sum of cations or sum of cations + sulfate. Must be in the same
		form as in the samdf dataframe. For example, if normalizing to the sum
		of cation charge, then:

			Sig_species = ['Ca_uM','K_uM','Mg_uM','Na_uM']

		If normalizing to the sum of cation + sulfate charge, then:

			Sig_species = ['Ca_uM','K_uM','Mg_uM','Na_uM', 'SO4_uM']

	n_iter : int
		The number of Monte Carlo iterations to perform. Defaults to ``10000''

	rmse_cutoff : float
		The RMSE cutoff to use when deciding which Monte Carlo iterations to
		store and which to discard. Defaults to ``0.01''.

	pyEMs: list
		List of letter codes for end-members to be included as pyrite weathering.
		If pyrite is treated as a separate end-member, then pyEMs = ['P'].
		Defaults to ``['P']''

	pptdf : pd.DataFrame
		DataFrame of precipitation (rainwater) chemistry used to correct sample
		data for atmospheric inputs. Must contain columns:

			Ca_uM, Cl_uM, K_uM, Mg_uM, Na_uM, SO4_uM

	Returns
	-------
	fracs : pd.DataFrame
		Resulting fractional contributions of each end-member to each sample.
		Index names are the same as ``samdf''; columns contain the:

			mean, median, std. dev., 25th percentile, and 75th percentile

		of stored Monte Carlo results for each fractional contribution.

	Apost : pd.DataFrame
		Resulting posterior end-member compositional matrix of each end-member.
		Index names are the same as ``emdf''; columns contain the:

			mean, median, std. dev., 25th percentile, and 75th percentile

		of stored Monte Carlo results for each end-member composition.
	'''

	#------------------------------------------#
	# STEP 0: CORRECT FOR PRECIPITATION INPUTS  #
	#------------------------------------------#

	vs_ppt = ['Ca_uM', 'K_uM', 'Mg_uM', 'Na_uM', 'SO4_uM'] #non-Cl species

	Rp = {} #define dictionary to hold precipitation ratios
	Rp['Cl_uM'] = pptdf['Cl_uM'].median()

	for v in vs_ppt:
		Rp[v+'Cl_uM'] = (pptdf[v]/pptdf['Cl_uM']).median()

	#correct Cl- data into a new column, preserving original Cl_uM
	samdf['Cl_uM*'] = samdf['Cl_uM'] - Rp['Cl_uM']
	samdf.loc[samdf['Cl_uM*'] < 0, 'Cl_uM*'] = 0 #force negative values to be zero

	#correct all other concentration data using ratios and original Cl_uM
	for v in vs_ppt:
		samdf[v] = [r[v] - Rp[v+'Cl_uM']*min([r['Cl_uM'], Rp['Cl_uM']]) for i, r in samdf.iterrows()]
		samdf.loc[samdf[v] < 0, v] = 0

	samdf['Cl_uM'] = samdf['Cl_uM*']
	
	#---------------------------------------------------#
	# STEP 1: SCALE DATA AND GENERATE SAMPLE (B) MATRIX #
	#---------------------------------------------------#

	#extract end-member and tracer lists
	ems = emdf.columns.str[0].unique()
	trs = emdf.index
	n_em = len(ems)
	n_tr = len(trs)
	n_sam = len(samdf)

	#calculate Chi values and denominator charges (Sig)
	vs = ['Ca_uM','Cl_uM','K_uM','Mg_uM','Na_uM','SO4_uM']
	Zs = {'Ca_uM':2,'Cl_uM':1,'K_uM':1,'Mg_uM':2,'Na_uM':1,'SO4_uM':2}
	wts = pd.Series(Zs) #make into series for weighted sum

	Sig = (samdf[Sig_species]*wts[Sig_species]).sum(axis=1)

	vs = ['Ca_uM','K_uM','Mg_uM','Na_uM','SO4_uM','Cl_uM']
	vnus = ['Ca','K','Mg','Na','S','Cl'] #drop units

	for v, vnu in zip(vs, vnus):
		samdf['Chi'+vnu] = Zs[v]*samdf[v]/Sig

	#make matrix
	B = samdf.loc[:,trs]
	B['unity'] = 1 #add unity constraint column
	Ba = B.values #make np.ndarray variable and designate with 'a'

	#------------------------------------#
	# STEP 2: GENERATE DESIGN (A) MATRIX #
	#------------------------------------#

	# For each iteration, randomly draw each end member composition value from a 
	# uniform distribution constrained by the inputted bounds, taking into 
	# account prescribed co-variance. Include the "sum to unity" constraint as 
	# the final row of the matrix ``A''.

	#first, store 2d matrix for each end member in a dict
	emMats = {}
	for e in ems:
		
		#make randomly generated dataframe for each end member
		X = pd.DataFrame(np.random.uniform(0,1,[n_iter, n_tr]),
						 columns = trs,
						 index = np.arange(0,n_iter)
						)
		
		#constrain random numbers to be between boundaries
		X = X*(emdf[e+'_max'] - emdf[e+'_min']) + emdf[e+'_min']
		
		#manually change co-varying tracers. See Torres et al. (2016)
		#dolomite vs. calcite but if you have separate calcite and dolomite
		if e == 'L':
			#Mg and Ca co-variance in limestone
			X['ChiMg'] = 1 - X['ChiCa']
		
		#halite vs. gypsum/anhydrite
		if e == 'E':
			#assume all Cl from halite and no other Na source. But this assumption won't hold for hot springs
			X['ChiCl'] = X['ChiNa']
			
			#assume Ca and Na as the only base cations
			X['ChiCa'] = 1 - X['ChiNa']

			#assume sulfur balances calcium
			X['ChiS'] = X['ChiCa']
		
		#add unity constraint column
		X['unity'] = 1
		
		#save to dictionary
		emMats[e] = X
		
	#next, stack into 3d np.ndarray of A matrices
	Aa = np.dstack([emMats[k].values for k in emMats.keys()])

	#--------------------------------------------------#
	# STEP 3: SOLVE NON-NEGATIVE LEAST SQUARES PROBLEM #
	#--------------------------------------------------#
	
	# For each iteration, find X, the matrix of fractional lithology 
	# contributions that best describes the data matrix B given the design 
	# matrix A for that iteration. That is, solve:
	# 
	# \min_{\mathbf{X}} || \mathbf{AX} - \mathbf{B} ||
	# 
	# subject to the constraint that each element in X is non-negative. This is 
	# equivalent to performing a simultaneous set of linear regressions for the 
	# equations:
	# 
	# $$\chi_{X^*}^{r} = \sum_{i = 1}^{n} f^i \chi_{X^*}^i$$
	# 
	# where subscript $r$ refers to measured river water and superscript $i$ (
	# $=$ D, E, G, L, or S) refers to each lithology end member. Then, calculate 
	# the root mean square error (RMSE) for each iteration following:
	# 
	# $$ RMSE = ( \frac{|| \mathbf{AX} - \mathbf{B} ||^2}{n_{sam}} )^{1/2} $$
	# 
	# where $n_{sam}$ is the number of samples in the sample set.

	#perform regressions
	res = [nnls(Aa[i,:,:],Ba[j,:])[0] for i,j in product(range(n_iter),range(n_sam))]

	#reshape results matrix
	Xa = np.array(res).reshape(n_iter,n_sam,-1)

	#calculate SSE and RMSE for each model iteration
	SSE = [(norm(np.dot(Aa[i,:,:],Xa[i,:,:].T) - Ba.T))**2 for i in range(n_iter)]
	SSEa = np.array(SSE)
	RMSEa = (SSEa/n_sam)**0.5

	#------------------------------------------------------#
	# STEP 4: EXTRACT LOW-RMSE SOLUTIONS AND STORE RESULTS #
	#------------------------------------------------------#

	#extract indices
	rmse_ind = np.argsort(RMSEa)[:int(rmse_cutoff*n_iter)]

	#calculate fractional contribution statistics
	f_mean = pd.DataFrame(np.mean(Xa[rmse_ind,:,:],axis=0),
						  index = samdf.index,
						  columns = ['f'+e+'_mean' for e in ems])

	f_std = pd.DataFrame(np.std(Xa[rmse_ind,:,:],axis=0),
						 index = samdf.index,
						 columns = ['f'+e+'_std' for e in ems])

	f_med = pd.DataFrame(np.median(Xa[rmse_ind,:,:],axis=0),
						 index = samdf.index,
						 columns = ['f'+e+'_med' for e in ems])

	f_25pctle = pd.DataFrame(np.quantile(Xa[rmse_ind,:,:],0.25,axis=0),
						 index = samdf.index,
						 columns = ['f'+e+'_25pctle' for e in ems])

	f_75pctle = pd.DataFrame(np.quantile(Xa[rmse_ind,:,:],0.75,axis=0),
						 index = samdf.index,
						 columns = ['f'+e+'_75pctle' for e in ems])

	#save fracs to dataframe
	fracs = pd.concat([f_mean,f_std,f_med,f_25pctle,f_75pctle],axis=1)

	#calculate end-member range statistics
	A_mean = pd.DataFrame(np.mean(Aa[rmse_ind,:-1,:],axis=0),
						  index = trs,
						  columns = [e+'_mean' for e in ems])

	A_std = pd.DataFrame(np.std(Aa[rmse_ind,:-1,:],axis=0),
						 index = trs,
						 columns = [e+'_std' for e in ems])

	A_med = pd.DataFrame(np.median(Aa[rmse_ind,:-1,:],axis=0),
						 index = trs,
						 columns = [e+'_med' for e in ems])

	A_25pctle = pd.DataFrame(np.quantile(Aa[rmse_ind,:-1,:],0.25,axis=0),
						 index = trs,
						 columns = [e+'_25pctle' for e in ems])

	A_75pctle = pd.DataFrame(np.quantile(Aa[rmse_ind,:-1,:],0.75,axis=0),
						 index = trs,
						 columns = [e+'_75pctle' for e in ems])

	#save posterior distributions to dataframe
	Apost = pd.concat([A_mean,A_std,A_med,A_25pctle,A_75pctle],axis=1)

	#-------------------------------#
	# STEP 5: SAVE OUTPUTS TO FILES #
	#-------------------------------#

	# Write all result DataFrames to CSV files in a fixed 'outputs' subfolder
	# relative to the current working directory. The folder is created if it
	# does not already exist.

	output_dir = 'outputs'
	os.makedirs(output_dir, exist_ok = True)
	B.to_csv(os.path.join(output_dir, 'B.csv'))
	fracs.to_csv(os.path.join(output_dir, 'fracs.csv'))
	Apost.to_csv(os.path.join(output_dir, 'Apost.csv'))

	return B, fracs, Apost