#################################
# Copyright 2020 OTICS Advanced Analytics Inc.
############################################

import numpy as np
import pandas as pd
import numpy as np
from scipy import stats
import numpy_indexed as npi
import sys
import math
from functools import reduce
from itertools import chain
import csv
from scipy.stats import kstest, norm
from statsmodels.stats.diagnostic import lilliefors
import random
import maads
import json
import ast
#from statsmodels.stats import shapiro
#from scipy import stats

#import plot

def getdata(idxall,df,cols,filename,headers):
    maindata=[]

    rows=0
    with open(filename, 'w') as myfile:
      # write headers
      rb=[] 
      for h in headers:
          rb.append(h)
      buf=",".join(rb)    
      myfile.write(buf + "\n")
      
          
      for idx in idxall:
        rb=[]      
        for c in range(cols):
          d=df.iloc[idx,c]
          rb.append(d)
       # print(rb)  
        buf=','.join(map(str, rb)) 
      #  print(buf)
        myfile.write(buf + "\n")
        rows+=1

      return rows

def balancebigdata(csvpath,numbins,maxrows,tofile,bincutoff,distcutoff,skipcols=0):
        binsarr=list()
        colind=0
        bucketlist=[]

        try:
          df=pd.read_csv(csvpath)
          df=df.drop_duplicates()
        except Exception as e:
          print("ERROR: reading file: %s" % e)
          return

        headers = df.columns
        bestvars=[]
        vardetails=[]
        countallocv=[]
        countsv=[]
        lmv=[]
        countalloc=[]
        binsv=[]

#        print(headers)

 #       print(len(df))
        jbuf=""
        for column in df:
          if colind>skipcols-1: # skip col=0   
              try:    
                cols=df[column].drop_duplicates().to_numpy()
                cnts, bins = np.histogram(cols, bins=numbins)
                bins=[round(b,3) for b in bins]
                ind = np.digitize(cols, bins)
                stat,bincut,distcut,mybins,share,cntsv,means,lm=countnumbuckets(ind,column,bins,bincutoff,numbins,cnts,distcutoff)
                #print("share:",share)
                sharearr=[round(s,3) for s in share]
              #  jbuf=column+":{'CountAllocation':"+str(sharearr) + "},{'Counts':"+str(cntsv)+"},{'Lilliefors-Statistics':"+str(round(lm[0],3))+'}'
                countallocv="CountAlloc:"+str(list(sharearr))
                countsv="BinCounts:"+str(list(cntsv))
                lmv="Lilliefors:"+str(round(lm[0],3))
                binsv="Bins:"+str(list(mybins))
                
                vardetails.append([column,countsv,countallocv,binsv,lmv]) 
                if stat==1:
                  d=equal_bin(cols,numbins,cnts,bins,ind)
                  bucketlist.append(d)
                  bestvars.append([column,bincut,distcut])
              except Exception as e:
                print("ERROR: reading columns: %s" % e)
                continue
              
          colind+=1
        try:
          idx,lil,mk=mergelists(bucketlist,numbins,maxrows)
        except Exception as e:
          print("ERROR: merging lists: %s" % e)
          return
        try:
        #  print(vardetails)  
          rc=getdata(mk,df,colind,tofile,headers)
       #   jbuf+="]"
          jsonstr=str({'Bins':list(bins), 'OptimizedCounts':list(idx),'OriginalCounts':list(cnts),'Lilliefors-Statistic':round(lil[0],3),'Numbins':numbins+1,
                       'Maxrows':maxrows,'Bincutoff':bincutoff,\
              'Distcutoff':distcutoff,'OriginalDataSize':len(df),'BalancedDataSize':rc,'BestVariables':bestvars,'InputFile':csvpath,'OutputFile':tofile,
                       'VariableDetails':vardetails})

          json_string = json.dumps(ast.literal_eval(jsonstr))
          print(json_string, "\nFile Output Complete")
        except Exception as e:
          print("ERROR: analysing data : %s" % e)
          return
        

# see which variables have the best distribution of buckets
def getlillie(share):
    return lilliefors(share)

def countnumbuckets(inds,colname,bins,bincutoff,numbins,cnts,distcutoff):
        np.set_printoptions(threshold=sys.maxsize)

     #   print(colname,inds)
        mybins = list(set(inds))
        thresh=len(mybins)/(numbins+1)
        share=cnts/sum(cnts)

        means=sum(share)/len(share)
        lm=getlillie(share)
        
        #print(mybins)
        #print(share)
        #print(cnts)

        #print(means)
        #print(lm)
        
        if thresh >= bincutoff and lm[0]<=distcutoff:
                return 1,round(thresh,3),round(lm[0],3),mybins,share,cnts,means,lm
        else:
             #   print("Bad bins:",mybins,colname,len(mybins),thresh,cnts,share,means)
                return 0,round(thresh,3),round(lm[0],3),mybins,share,cnts,means,lm


def getrowsfrombuckets(drows,mpall):
        a=1
        mk=[]
        #print(len(drows),len(mpall))
        for r,m in zip(drows,mpall):
             #print(m)   
             mk+=random.sample(m, r)
             #mk.append(val)
        ms=set(mk)     
        return ms
                
        
def equal_bin(data, numbins,cnts,bins,out):
   np.set_printoptions(threshold=sys.maxsize)   
    # parameter q specifies the number of bins
#   out,bins = pd.cut(data, q, labels=False, retbins=True, right=False)
#   print(out) # contains bin indices for data
#   print("Orginal:",obins,cnt/len(out))
   totdata=[]
  # totdatall=[]
   bucketindexes=[]
   arrindices=[]
   d=[]


   #get indices of data in buckets
   k=0            
   for i in out:
       bucketindexes.append([k,i])
       k+=1

   #print(bucketindexes)
   d=[]
   d2=[]
   bc=1# bucket count
   for nb in range(numbins+1):
     totitems=0
     barr=[]
    # print("NB:",nb,bc)
     for ir in bucketindexes:
#        if totitems<nb and bc==ir[1]:
#        print(ir)
        if bc==ir[1]:
          barr.append(ir[0])    
          totitems+=1 
     bc+=1
     d.append(barr)

 #  print("BARR DATA:",bc)
#   print(len(d))
  
   return d  

def merge(lsts):
    sets = [set(lst) for lst in lsts if lst]
    merged = True
    while merged:
        merged = False
        results = []
        while sets:
            common, rest = sets[0], sets[1:]
            sets = []
            for x in rest:
                if x.isdisjoint(common):
                    sets.append(x)
                else:
                    merged = True
                    common |= x
            results.append(common)
        sets = results
    return list(sets[0])

def mergelists(dall,numbins,maxrows):
   mainlist=[]
   mlist=[]
   mainmp=[]
   mpall=[]
   mpallset=[]
   dlens=[]
   dlenshare=[]
   drows=[]
   for i in range(numbins+1):
     d=[]      
#     mpallset=[]
#     dlens=[]
     dlenshare=[]
     for di in dall: #list of variables-each variable has buckets
         #d.append(set(di[i]))
         d+=set(di[i])  #merge it
     mp=set(d)
     mpall+=mp
     mpallset.append(mp)
     dlens.append(len(mp))
     mainmp.append(mp)
   dlenshare=[d/sum(dlens) for d in dlens]
   drows=[int(dr*maxrows) for dr in dlenshare ]
   mk=getrowsfrombuckets(drows,mpallset)
   sharemk=[m/sum(mk) for m in mk]
   lil=getlillie(sharemk)
   
   
   return drows,lil,mk
        
def intersect(dall,numbins,maxrows):

#   len(dall)
   mainlist=[]
   for i in range(numbins+1):
     d=[]      
     for di in dall: #list of variables-each variable has buckets
         d.append(di[i])

     mp=list(set(chain.from_iterable(d)))
     mainlist.append(mp)
   mk=list(reduce(set.intersection, [set(item) for item in mainlist ]))

   return mk


#csvp="C:\MAADS\Companies\Precision Drilling\PoC\Data\decision2_cleanholetime\cleanholetimeD2training.csv"
#tofile1="C:\MAADS\Companies\Precision Drilling\PoC\Data\decision2_cleanholetime\cleanholetimeD2trainingoptimized.csv"

#csvp="C:\MAADS\Companies\Precision Drilling\PoC\Data\decision2_cleanholetime\cleanholeD2training2.csv"
#csvp="C:\\MAADS\\Companies\\Precision Drilling\\PoC\\Data\\decision3_backream\\backreamstandD3training.csv"
#tofile="C:\\MAADS\\Companies\\Precision Drilling\\PoC\\Data\\decision3_backream\\backreamstandD3trainingoptimized.csv"

csvp2="C:\\MAADS\\Companies\\Precision Drilling\\PoC\\Data\\decision4a_stickslip\\backreamD4Atraining.csv"
tofile2="C:\\MAADS\\Companies\\Precision Drilling\\PoC\\Data\\decision4a_stickslip\\backreamD4Atrainingoptimized.csv"

bincutoff=.7
distcutoff=.5
numbins=10
#balancebigdata(csvp2,numbins,2000,tofile2,bincutoff,distcutoff,1)


#maads.balancebigdata(csvp2,numbins,2000,tofile2,bincutoff,distcutoff,1)



