import sys,os, pickle, numpy, pylab, operator, itertools
import cv2
from shutil import copy as copyfile
from PyQt4.QtCore import *
from PyQt4.QtGui import *
import matplotlib.pyplot as plt
from DataParseApp import dataparseDialog
from sklearn.decomposition import NMF

projectpath=os.path.split(os.path.abspath(__file__))[0]
sys.path.append(os.path.join(projectpath,'ui'))

pythoncodepath=os.path.split(projectpath)[0]
jcapdataprocesspath=os.path.join(pythoncodepath, 'JCAPDataProcess')
sys.path.append(jcapdataprocesspath)
from VisualizeDataApp import visdataDialog

sys.path.append(os.path.join(jcapdataprocesspath,'AuxPrograms'))
from fcns_ui import *
from fcns_io import *

platemapvisprocesspath=os.path.join(pythoncodepath, 'JCAPPlatemapVisualize')
sys.path.append(platemapvisprocesspath)
from plate_image_align_Dialog import plateimagealignDialog

import numpy as np

###############UPDATE THIS TO BE THE FOLDER CONTAINING parameters.py
#paramsfolder=r'K:\users\hte\Raman\39664\20170607analysis_tippy-top'
paramsfolder=r'K:\users\hte\Raman\33444\20170608analysis'

#if not paramsfolder is None:
sys.path.append(paramsfolder)
from parameters import *
#else:
#    plateidstr='3344'
#
#    pathd={'ramanfile':r'K:\users\hte\Raman\33444\HSS_33444_map-1-_CRR-EM-copy.txt'}
#    pathd['mainfolder']=os.path.split(pathd['ramanfile'])[0]
#    pathd['savefolder']=os.path.join(pathd['mainfolder'], '20170607analysis')
#    pathd['infopck']=pathd['ramanfile'][:-4]+'__info.pck'
#    pathd['allspectra']=os.path.join(pathd['savefolder'],'allspectra.npy')
#    pathd['nmfdata']=os.path.join(pathd['savefolder'],'nmf4.pck')
#    pathd['edges']=os.path.join(pathd['savefolder'],'edges.png')
#    pathd['mapfill']=os.path.join(pathd['savefolder'],'blobmap.png')
#    pathd['blobd']=os.path.join(pathd['savefolder'],'blobd.pck')
#    pathd['alignedsamples']=os.path.join(pathd['savefolder'],'alignedsamples.png')
#    pathd['alignedsamplestxt']=os.path.join(pathd['savefolder'],'alignedsamples.txt')
#    pathd['spectrafolder']=os.path.join(pathd['savefolder'],'sample_spectra')
#    pathd['map']=os.path.join(pathd['spectrafolder'],'raman_sample_index_map.map')
#    pathd['samplepixels']=os.path.join(pathd['spectrafolder'],'samplepixels.png')
#    pathd['udibasepath']=os.path.join(pathd['savefolder'],'ave_rmn_')
#
#    udi_ternary_projection_inds=[0, 1, 2]#only used for the all.udi file
#
#    sample_list=[1850,1851,1852,1853,1854,1855,1905,1906,1907,1908,1909,1910,1911,1912,1913,1914,1915,1916,1917,1918,1919,1969,1970,1971,1972,1973,1974,1975,1976,1977,1978,1979,1980,1981,1982,1983,2033,2034,2035,2036,2037,2038,2039,2040,2041,2042,2043,2044,2045,2046,2047,2097,2098,2099,2100,2101,2102,2103,2104,2105,2106,2107,2108,2109,2110,2111]
#    dx_smp=1.
#    dy_smp=1.
#
#    default_sample_blob_dict=dict({}, \
#    smp_is_square=0, smp_width=1., bcknd_is_square=0,  bcknd_min_width=1.3, bcknd_max_width=1.4, removedups=1\
#    )
#
#    show_help_messages=True




platemappath=getplatemappath_plateid(plateidstr)

if not os.path.isdir(pathd['mainfolder']):
    print 'NOT A VALID FOLDER'
if not os.path.isdir(pathd['savefolder']):
    os.mkdir(pathd['savefolder'])
if not os.path.isdir(pathd['spectrafolder']):
    os.mkdir(pathd['spectrafolder'])
    
class MainMenu(QMainWindow):
    def __init__(self, previousmm, execute=True, **kwargs):
        super(MainMenu, self).__init__(None)
        self.parseui=dataparseDialog(self, title='Visualize ANA, EXP, RUN data', **kwargs)
        self.alignui=plateimagealignDialog(self, manual_image_init_bool=False)
        if execute:
            self.parseui.exec_()


def doNMF(datan,n_components=4):
    # from Mitsu
    #alternatively PCA ... might me faster
    nmf=NMF(n_components=n_components,init='nndsvd')
    data_decomp_all=nmf.fit_transform(datan)
    data_components_all=nmf.components_
    return data_decomp_all,data_components_all

def rgb_comp(arr2d, affine=True):
    cmy_cmyk=lambda a:a[:3]*(1.-a[3])+a[3]
    rgb_cmy=lambda a:1.-a
    rgb_cmyk=lambda a:rgb_cmy(cmy_cmyk(a))


    return numpy.array([rgb_cmyk(a) for a in arr2d])
        
def imGen(data_decomp_all,ramaninfod,cmykindeces=[3, 2, 1, 0]):
    cmykvals=copy.copy(data_decomp_all[:, cmykindeces])
    cmykvals/=cmykvals.max(axis=0)[numpy.newaxis, :]
    img=numpy.reshape(rgb_comp(cmykvals), (ramaninfod['xshape'], ramaninfod['yshape'], 3))
    return img

def findEdges(img_gray):
    #this uses automatic thresholding from one of the cv2 tutorials
    sigma = 0.33
    v = np.median(img_gray)
    lower = int(max(0, (1.0 - sigma) * v))
    upper = int(min(255, (1.0 + sigma) * v))
    edges = cv2.Canny(np.uint8(img_gray),lower,upper)
    return edges

def findContours(edges):
    #the contours are now found by searching the most external convex hull
    #this way mos of the not fully closed samples are detected as well
    im2, contours, hierarchy = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
    iWithContour = cv2.drawContours(edges, contours, -1, (255,20,100), 5)
    mapimage = np.zeros_like(edges)
    #this fills the contours
    for i in range(len(contours)):
        cv2.drawContours(mapimage, contours, i, color=255, thickness=-1)
    #this is to calculate the center of each contour
    x=[]
    y=[]
    for c in contours:
        # compute the center of the contour
        M = cv2.moments(c)
        try:
            x.append(M['m10']/(M['m00']))
            y.append(M['m01']/(M['m00']))
        except:
            #this was nessesary as the divisor is sometimes 0
            #yield good results but should be done with caution
            x.append(M['m10']/(M['m00']+1e-23))
            y.append(M['m01']/(M['m00']+1e-23))
    return iWithContour, mapimage, contours, x, y
    


mainapp=QApplication(sys.argv)
form=MainMenu(None, execute=False)
#form.show()
#form.setFocus()
#mainapp.exec_()

parseui=form.parseui
alignui=form.alignui
parseui.rawpathLineEdit.setText(pathd['ramanfile'])
parseui.infopathLineEdit.setText(pathd['infopck'])
parseui.getinfo(ramaninfop=pathd['infopck'], ramanfp=pathd['ramanfile'])#opens or creates

if os.path.isfile(pathd['allspectra']):
    with open(pathd['allspectra'], mode='rb') as f:
        fullramandataarray=numpy.load(f)
elif 1:
    fullramandataarray=parseui.readfullramanarray(pathd['ramanfile'])#opens or creates
    with open(pathd['allspectra'], mode='wb') as f:
        numpy.save(f, fullramandataarray)
    
ramaninfod=parseui.ramaninfod
#parseui.exec_()
#ramaninfod['number of spectra']
#ramaninfod['xdata']
#ramaninfod['ydata']
#ramaninfod['Wavenumbers_str']
#ramaninfod['Spectrum 0 index']
ramaninfod['xdata']/=1000.
ramaninfod['ydata']/=1000.#convert to mm
ramaninfod['xshape']= len(np.unique(ramaninfod['xdata']))
ramaninfod['yshape']= len(np.unique(ramaninfod['ydata']))

ramaninfod['dx']= (ramaninfod['xdata'].max()-ramaninfod['xdata'].min())/(ramaninfod['xshape']-1)
ramaninfod['dy']= (ramaninfod['ydata'].max()-ramaninfod['ydata'].min())/(ramaninfod['yshape']-1)

nx=dx_smp/ramaninfod['dx']
ny=dy_smp/ramaninfod['dy']
ntot=nx*ny

ramanreshape=lambda arr: np.reshape(arr, (ramaninfod['xshape'], ramaninfod['yshape'])).T[::-1, ::-1]

ramannewshape=(ramaninfod['yshape'], ramaninfod['xshape'])

image_of_x=ramanreshape(ramaninfod['xdata'])
image_of_y=ramanreshape(ramaninfod['ydata'])
    
    
#extent=[ramaninfod['xdata'].max(), ramaninfod['xdata'].min(), ramaninfod['ydata'].min(), ramaninfod['ydata'].max()]
#extent=[ramaninfod['xdata'].max(), ramaninfod['xdata'].min(), ramaninfod['ydata'].max(), ramaninfod['ydata'].min()]
extent=[image_of_x[0, 0], image_of_x[-1, -1], image_of_y[0, 0], image_of_y[-1, -1]]

def ramanimshow(im, **kwargs):
    plt.imshow(im, origin='lower', interpolation='none', aspect=1, extent=extent, **kwargs)

if os.path.isfile(pathd['nmfdata']):
    with open(pathd['nmfdata'], mode='rb') as f:
        tempd=pickle.load(f)
    data_decomp_all,data_components_all,rgbimagedata=[tempd[k] for k in 'data_decomp_all,data_components_all,rgbimagedata'.split(',')]
else:
    data_decomp_all,data_components_all = doNMF(fullramandataarray,4)
    #rgbimagedata=imGen(data_decomp_all,ramaninfod)
    rgbimagedata=np.zeros(ramannewshape+(3,), dtype='float32')
    for i, arr in enumerate(data_decomp_all[:, :3].T):
        rgbimagedata[:, :, i]=np.array([ramanreshape(arr/arr.max())])
    tempd={}
    tempd['data_decomp_all']=data_decomp_all
    tempd['data_components_all']=data_components_all
    tempd['rgbimagedata']=rgbimagedata
    with open(pathd['nmfdata'], mode='wb') as f:
        tempd=pickle.dump(tempd, f)



if 1 and os.path.isfile(pathd['blobd']):
    with open(pathd['blobd'], mode='rb') as f:
        blobd=pickle.load(f)
else:
    edges = np.zeros(ramannewshape, dtype='uint8')
    plt.clf()
    for i in range(data_decomp_all.shape[1]):
        arr=np.uint8(ramanreshape(data_decomp_all[:,i])/data_decomp_all[:,i].max()*254)
        edgetemp=findEdges(arr)
    #    plt.imshow(edgetemp)
    #    plt.show()
        edges[np.where(edgetemp>0)] = 244

    ramanimshow(edges)
    plt.savefig(pathd['edges'])
    plt.clf()

    im2, contours, hierarchy = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)

    image_of_inds=ramanreshape(numpy.arange(ramaninfod['number of spectra']))
    

    mapfill = np.zeros(ramannewshape, dtype='uint8')