import os, sys
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
#from creader import ClusterFileReader
import creader as cr
#import clusterFunctions as cf
from PIL import Image
import boost_histogram as bh

def color_images(r, emin, emax, ebins, xmin, xmax, ymin, ymax,  images=None, csize=3,gain=150):
    n=0
    while (((cl:=r.read(100000,None)).size) ):
        v=cr.clusterize(csize,cl['data'])
        vv=[cl['x'],cl['y'],v['tot']/gain]
        image,bins=np.histogramdd(vv, bins=[xmax-xmin,ymax-ymin,ebins], range=[[xmin,xmax-1],[ymin,ymax-1],[emin,emax]])

        if images is None:
            images = image.copy()  
        else:
            images=images+image
        n+=v.shape[0]
        #print(v.shape)
        print("Read ",n," clusters",np.sum(images))
        
    return images, bins

def color_images_bh(r, emin, emax, ebins, xmin, xmax, ymin, ymax,  hist3d=None, csize=3,gain=150,noisemap=None):
    if hist3d is None:
        hist3d = bh.Histogram(
            bh.axis.Regular(xmax-xmin+1, xmin-0.5, xmax+0.5),
            bh.axis.Regular(ymax-ymin+1, xmin-0.5, xmax+0.5),
            bh.axis.Regular(ebins, emin, emax),
        )
    n=0
    while (((cl:=r.read(100000,noisemap)).size)):#and (n<1E7)):
        v=cr.clusterize(csize,cl['data'])
        vv=[cl['x'],cl['y'],v['tot']/gain]
        hist3d.fill(*vv)
        n+=v.shape[0]
        if n%1000000 ==0:
            print("Read ",n," clusters",hist3d.sum())
    return hist3d
        
def analyze_clusters(r, emin, emax, ecutmin, ecutmax, xmin, xmax, ymin, ymax,   ietax=None, ietay=None, im=None, sp=None, etas=None, intim=None, csize=3,gain=150, nbins=100, etabins=250, subpix=5):
    if csize==3:
        etamin=-0.6
        etamax=0.6
    else:
        etamin=-0.1
        etamax=1.1
    ttx=None
    tty=None
    n=0
    while (cl:=r.read(100000,None)).size:
        v=cr.clusterize(csize,cl['data'])
        spectrum, ebins =np.histogram(v['tot'][np.where((cl['x']>=xmin) & (cl['x']<xmax) & (cl['y']>=ymin) & (cl['y']<ymax))]/gain, bins=nbins, range=[emin,emax], density=None, weights=None)
        image,xedges,yedges=np.histogram2d(cl['x'][np.where((v['tot']/gain>ecutmin) & (v['tot']/gain<ecutmax))],cl['y'][np.where((v['tot']/gain>ecutmin) & (v['tot']/gain<ecutmax))],bins=[xmax-xmin,ymax-ymin],range=[[xmin,xmax-1],[ymin,ymax-1]])
        eta,etabinsx,etabinsy=np.histogram2d(v['etax'][np.where((v['tot']/gain>ecutmin) & (v['tot']/gain<ecutmax) & (cl['x']>=xmin) & (cl['x']<xmax) & (cl['y']>=ymin) & (cl['y']<ymax))],v['etay'][np.where((v['tot']/gain>ecutmin) & (v['tot']/gain<ecutmax) & (cl['x']>=xmin) & (cl['x']<xmax) & (cl['y']>=ymin) & (cl['y']<ymax))],bins=[etabins,etabins],range=[[etamin,etamax],[etamin,etamax]])
        if im is None:
            im = image.copy()  
        else:
            im=im+image
        if sp is None:
            sp = spectrum.copy() 
        else:
            sp=sp+spectrum
        if etas is None:
            etas = eta.copy()
        else:
            etas=etas+eta
        if ietax is not None and ietay is not None:
            i=100
            if subpix!=2:
                ibx=np.searchsorted(etabinsx,v['etax'])
                iby=np.searchsorted(etabinsy,v['etay'])
            
                ibx[np.where(ibx>=etabinsx.shape[0]-1)]=etabinsx.shape[0]-2
                iby[np.where(iby>=etabinsy.shape[0]-1)]=etabinsy.shape[0]-2
                if csize==3:
                    px=cl['x']+ietax[ibx,iby]-0.5
                    py=cl['y']+ietay[ibx,iby]-0.5
                    #print("***",v['corner'][i],"\n",v['etax'][i],v['etay'][i],"\n",etabinsx[ibx[i]],etabinsy[iby[i]],"\n",ietax[ibx,iby][i],ietay[ibx,iby][i],"\n",cl['x'][i],cl['y'][i],"\n",px[i],py[i])
                else:
                    offx=v['corner']%2
                    offy=(v['corner']/2).astype(int)
                    px=cl['x'].astype(float)+(-1+offx.astype(float))+ietax[ibx,iby]
                    py=cl['y'].astype(float)+(-1+offy.astype(float))+ietay[ibx,iby]
                    #print("***",v['corner'][i],"\n",offx[i], offy[i],"\n",v['etax'][i],v['etay'][i],"\n",etabinsx[ibx[i]],etabinsy[iby[i]],"\n",ietax[ibx,iby][i],ietay[ibx,iby][i],"\n",cl['x'][i],cl['y'][i],"\n",px[i],py[i],"\n",((-1+offx.astype(float))+ietax[ibx,iby])[i],((-1+offy.astype(float)))[i]+ietax[ibx,iby][i])
            else:
                offx=v['corner']%2
                offy=(v['corner']/2).astype(int)
                px=cl['x'].astype(float)+(0.25+0.5*offx.astype(float))
                py=cl['y'].astype(float)+(0.25+0.5*offy.astype(float))
                #print(v['corner'][i],offx[i], offy[i],v['etax'][i],v['etay'][i],cl['x'][i],cl['y'][i],px[i],py[i])
            intimage,xedges,yedges=np.histogram2d(px[np.where((v['tot']/gain>ecutmin) & (v['tot']/gain<ecutmax))],py[np.where((v['tot']/gain>ecutmin) & (v['tot']/gain<ecutmax))],bins=[subpix*(xmax-xmin),subpix*(ymax-ymin)],range=[[xmin,xmax-1],[ymin,ymax-1]])
            if intim is None:
                print("new")
                intim = intimage.copy()  
            else:
                intim=intim+intimage


    if ietax is None or ietay is None:
        return im, sp, ebins, etas, etabinsx, etabinsy
    else:
        return im, intim, sp, ebins, etas, etabinsx, etabinsy



def make_eta(r, emin, emax, ecutmin, ecutmax, xmin, xmax, ymin, ymax, im=None, sp=None, etas=None, csize=3,  gain=150, nbins=100, etabins=250):
    return analyze_clusters(r, emin, emax, ecutmin, ecutmax, xmin, xmax, ymin, ymax,   None, None, im, sp, etas, None, csize, gain, nbins, etabins)

def interpolate(r, emin, emax, ecutmin, ecutmax, xmin, xmax, ymin, ymax, ietax, ietay, im=None, sp=None, etas=None, intim=None, csize=3,gain=150, nbins=100, etabins=250):
    
    return analyze_clusters(r, emin, emax, ecutmin, ecutmax, xmin, xmax, ymin, ymax,   ietax, ietay, im, sp, etas, intim, csize, gain, nbins, etabins)


def prepare_interpolation(eta):
    ietax=np.cumsum(eta,axis=0)
    ietay=np.cumsum(eta,axis=1)
    netax=np.tile(ietax[-1,:],ietax.shape[0]).reshape(ietax.shape)
    netax[np.where(netax==0)]=1
    ietax=ietax/netax
    netay=np.transpose(np.tile(ietay[:,-1],ietay.shape[1]).reshape(ietay.shape))
    netax[np.where(netay==0)]=1
    ietay=ietay/netay
    return ietax, ietay


def interpolate_cl(etax,etay,ietax,ietay, etaxbins, etaybins):
    ibx=(np.abs(etabinsx - etax)).argmin()
    iby=(np.abs(etabinsy - etay)).argmin()
    px=ietax[ibx,iby]
    py=ietay[ibx,iby]
    return px, py


def plot_colz(hist2d, vmax=-1, vmin=0):
    if vmax<=0:
        vmax=np.max(hist2d)
    if vmin>vmax:
        vmin=0
    fig1, axs1 = plt.subplots()
    vv=axs1.imshow(np.transpose(hist2d),origin='upper',cmap=plt.cm.gray,vmax=vmax,vmin=vmin, interpolation='none' )
    fig1.colorbar(vv, ax=axs1)
    fig1.show()

def read_tiff_map(fname):
    im_tot = Image.open(fname)
    return np.transpose(np.flip(np.array(im_tot),axis=0))

def plothist2d(h):
    return plt.pcolormesh(*h.axes.edges.T, h.values().T)