#!/usr/bin/env python3
'''Rank self.correlation processing for a csv tabulation of self.counts by segment 
   First column is for whole crawl, then 100 columns for segs 0-99
   Each row is self.counts for some property, e.g. mime-detected or tld

   For example, assuming ALL.tsv has the whole-crawl warc-only self.counts
   and s...tsv have the segment self.counts, ALL with self.counts in column 1,

   tr -d ',' <ALL.tsv |head -100 | while read n m; do printf "%s%s\n" $n $(for i in {0..99}; do printf ",%s" $({ grep -w "w    $m\$" s${i}.tsv || echo NaN ;} | cut -f 1 ) ; done ) ; done > all_100.csv

   will produce such a file with
     * 100 rows, one for each of the top 100 self.counts
     * 101 columns, 0 for ALL and 1--100 for segs 0--99

   Usage: python3 -i spearman.py name id
     where name.csv has the input
'''
     
import numpy as np
from numpy import loadtxt
import scipy
from scipy import stats
import statsmodels.api as sm
import matplotlib.pyplot as plt
import pylab
from scipy.stats import _bootstrap

import sys, math, os.path, random

RNG = np.random.default_rng()
try:
  SS
except NameError:
  SS = {}

class CCStats:
  def __init__(self, year_month, name, file_base):
    global SS
    self.year_month = year_month
    self.name = name
    self.id = "%s %s"%(year_month,name)
    SS[self.id] = self
    self.file = os.path.expanduser("~/results/CC-MAIN-%s/%s.csv")%(year_month,
                                                                   file_base)
    self.counts=loadtxt(self.file,delimiter=',')
    self.N=self.counts.shape[1]-1
    # "If axis=0 (default), then each column represents a variable, with
    #        observations in the rows"
    # So each column is a sequence of self.counts, for whole crawl in column 0
    #   and for segments 0--self.N-1 in columns 1--self.N
    self.corr=stats.spearmanr(self.counts,nan_policy='omit').correlation

    self.all=self.corr[0][1:]
    self.all_s=stats.describe(self.all)
    self.all_m=self.all_s.mean
    self.all_sd=math.sqrt(self.all_s.variance)

    self.x=np.array([np.concatenate((self.corr[i][1:i],
                                self.corr[i][i+1:])) for i in range(1,self.N+1)])
    # The above, although transposed, works because the correlation matrix
    #  is symmetric
    self.xd=[stats.describe(self.x[i]) for i in range(self.N)]
    self.xs=stats.describe(np.array([self.xd[i].mean for i in range(self.N)]))
    self.xm=self.xs.mean
    self.xsd=np.sqrt(self.xs.variance)

    self.x_ranks=[stats.rankdata(-self.counts[:,i],method='average') for i in range(1,self.N+1)]

    self.aa = self.ranks()
    self.aa_by_all=self.aa[self.aa[:,1].argsort()]
    self.iranks_by_all = [int(r) for r in self.aa_by_all[:,0]]

  def qqa(self):
    # q-q plot for the whole crawl
    sm.qqplot(self.all, line='s')
    plt.gca().set_title('Rank self.correlation per segment wrt whole archive %s'%self.id)
    plt.show()

  def qqs(self):
    # q-q plots for the best and worst (by variance) segments
    global xv, xworst, xbest
    xv=[d.variance for d in self.xd]
    xworst=xv.index(max(xv))
    xbest=xv.index(min(xv))
    print(xbest,xworst)
    sm.qqplot(self.x[xbest], line='s')
    plt.gca().set_title('Best segment (least variance): %s'%xbest)
    plt.show()
    sm.qqplot(self.x[xworst], line='s')
    plt.gca().set_title('Worst segment (most variance): %s'%xworst)
    plt.show()

  def plot_x(self,sort=False,block=True,all_only=True,title=None):
    # Make these two subplots, w. and w/o sorting
    # See https://stackoverflow.com/questions/4700614/how-to-put-the-legend-outside-the-plot
    #  for legend hacking
    if sort:
      aso=np.argsort(-self.all)
      plot_all=self.all[aso]
      plot_x=np.array([self.xd[i].mean for i in range(self.N)])[aso]
    else:
      plot_all=self.all
      plot_x=[self.xd[i].mean for i in range(self.N)]
    if title is None:
      l1='Rank self.correlation of segment self.x whole crawl'
      l2='Mean of segment self.x whole crawl'
      plt.legend(loc='best',fontsize='small')
    else:
      l1=l2=None
    plt.plot(plot_all,'rx',label=l1)
    plt.plot([0,self.N-1],[self.all_m,self.all_m],'r',label=l2)
    if not(all_only):
      plt.plot(plot_x,'bx',label='Mean of rank self.correlation of each segment self.x self.all other segments')
      plt.plot([0,self.N-1],[self.xm,self.xm],'b',label='Mean of segment self.x segment means')
    plt.axis([0,self.N-1,0.85 if all_only else 0.8,1.0])
    plt.grid(True)
    if title is not None:
      plt.title(title)
    plt.show(block=block)

  def hist_x(self,align='mid'):
    hist(self.xm,self.xsd,[self.xd[i].mean for i in range(self.N)],
         'Mean of rank self.correlation of each segment self.x self.all other segments',
         align)

  def hist_all(self,align='mid'):
    hist(self.all_m,np.sqrt(self.all_s.variance),self.all,
         'Rank self.correlation of each segment self.x whole crawl %s'%self.id,
         align)

  def plot_ci(self,rhos=None,bootstrap=None,n=None,trim=None,conf=0.95,block=True):
    global cc
    # rhos are (rank) correlation values
    # If bootstrap is present, it should be the data from which rhos
    #  was derived via spearmanr, which will be used for bootstrapping, e.g. self.counts
    if rhos is None:
      rhos = self.all
    if n is None:
      n = rhos.shape[0]
    all_order=(-rhos).argsort()
    rhos_s=rhos[all_order]
    if trim is None:
      l=len(rhos)
    else:
      rhos_s=rhos_s[:trim]
      l=trim
    if bootstrap is not None:
      all=bootstrap[:,0] # [all_order]
      rest=bootstrap[:,1:] # [all_order]
      cc=np.array([tuple(stats.bootstrap((all,
                                          rest[:,i]),
                                          sr, vectorized=bool(print(i)),
                                          paired=True, method='bca',
                                          random_state=RNG).confidence_interval)
                    for i in range(l)])
      #cc=cc[(-rhos).argsort()].T
      cc=cc.T
      cc=[cc[0][all_order],cc[1][all_order]]
    else:
      cc=np.array([ci(r,n,conf) for r in rhos_s]).T
    ue=cc[1]-rhos_s
    le=rhos_s-cc[0]
    #for i in range(len(rhos)):
      #print(cc[i][0],rhos_s[i]-cc[i][0],rhos_s[i],cc[i][1],-rhos_s[i]+cc[i][1])
    plt.errorbar(np.arange(l),rhos_s,yerr=[le,ue],fmt='o')
    plt.title("Rank self.correlation of segments self.x whole archive %s\nwith confidence bars at %d%%"%(self.id,conf*100))
    plt.show(block=block)

  def ranks(self):
    # Combine segment measures:
    #  segID,rank self.corr. wrt self.all,inverse variance, mean cross rank self.corr.,first disagreement
    # convert to ranks, smallest value == highest rank
    all_ranked=stats.rankdata(-self.all,method='average') # invert since
                                                     #  large self.corr is good
    x_variance_ranked=stats.rankdata([self.xd[i].variance for i in range(self.N)])
                                                    # small self.corr variance is good
    x_mean_ranked=stats.rankdata([-(self.xd[i].mean) for i in range(self.N)])
                                                     # invert since
                                                     #  large mean self.corr is good
    fd_ranked=stats.rankdata([-first_diff(self.x_ranks[i]) for i in range(self.N)])
                                                     # invert since
                                                     #  large first diff is good
    return np.array([[i,
                      all_ranked[i],
                      x_variance_ranked[i],
                      x_mean_ranked[i],
                      fd_ranked[i]] for i in range(self.N)])

def hist(m,sd,hh,title,align):
  sdd=[(i,m-(i*sd)) for i in range(-2,3)]
  fig,hax=plt.subplots() # Thanks to https://stackoverflow.com/a/7769497
  sdax=hax.twiny()
  hax.hist(hh,color='lightblue',align=align)
  hax.set_title(title)
  for s,v in sdd:
       sdax.plot([v,v],[0,18],'b')
  sdax.set_xlim(hax.get_xlim())
  sdax.set_ylim(hax.get_ylim())
  sdax.set_xticks([v for s,v in sdd])
  sdax.set_xticklabels([str(s) for s,v in sdd])
  plt.show()

def first_diff(ranks):
  # first disagreement with baseline == {1,2,...}
  for i in range(len(ranks)):
    if ranks[i]!=i+1.0:
      return i
  return i+1

def ci(rho,n,conf=0.95):
  # Courtesy of https://stats.stackexchange.com/a/18904
  # rho is (rank) correlation, n is sample size
  stderr=1.0/math.sqrt(n-3)
  z=stats.norm.ppf(1.0-((1.0-conf)/2))
  delta=z*stderr
  lower=math.tanh(math.atanh(rho)-delta)
  upper=math.tanh(math.atanh(rho)+delta)
  return (lower,upper)

def sr(x,y):
  return stats.spearmanr(x,y)[0]

def _percentile_along_axis(theta_hat_b, alpha):
    """`np.percentile` with different percentile for each slice."""
    # the difference between _percentile_along_axis and np.percentile is that
    # np.percentile gets _all_ the qs for each axis slice, whereas
    # _percentile_along_axis gets the q corresponding with each axis slice
    shape = theta_hat_b.shape[:-1]
    alpha = np.broadcast_to(alpha, shape)
    percentiles = np.zeros_like(alpha, dtype=np.float64)
    for indices, alpha_i in np.ndenumerate(alpha):
        if np.isnan(alpha_i):
            # See https://github.com/scipy/scipy/pull/14351/commits/e5f38dee638d63981e89b4630e2ea5bc3a59a51d
            percentiles[indices] = np.nan
        else:
            theta_hat_b_i = theta_hat_b[indices]
            percentiles[indices] = np.percentile(theta_hat_b_i, alpha_i)
    return percentiles[()]  # return scalar instead of 0d array

_bootstrap._percentile_along_axis = _percentile_along_axis

def explore_deltas(basis_ym, basis_names, target_ym, target_names, nn):
  for n in nn:
    for p in basis_names:
      for q in target_names:
        basis=SS[f'{basis_ym} {p}']
        target=SS[f'{target_ym} {q}']
        m = int((basis.N - n) / 2)
        best = stats.describe(target.all[np.array(basis.iranks_by_all[:n])]).mean
        middle = stats.describe(target.all[np.array(basis.iranks_by_all[m:m+n])]).mean
        worst = stats.describe(target.all[np.array(basis.iranks_by_all[-n:])]).mean
        delta1 = best-worst
        stdev = math.sqrt(target.all_s.variance)
        ds1 = delta1 / stdev
        delta2 = abs(middle-target.all_m)
        ds2 = delta2 / stdev
        print(n,p,q,'%6.4g'%delta1,'%4.2g'%ds1,'%6.4g'%delta2,'%4.2g'%ds2)

def heat_map(basis_ym, basis_names, target_ym, target_names, nn, outf = sys.stdout):
  nn=nn if isinstance(nn,list) else list(nn)
  with (outf if outf is sys.stdout else open(outf,'w')) as outfile:
    print("\tN",file=outfile)
    print("Predict","\t".join(["{:d}"]*len(nn)).format(*nn),sep="\t",file=outfile)
    for p in basis_names:
      for q in target_names:
        basis=SS[f'{basis_ym} {p}']
        target=SS[f'{target_ym} {q}']
        tmin,tmax = target.all_s.minmax
        twidth = tmax - tmin
        best = (stats.describe(target.all[np.array(basis.iranks_by_all[:n])]).mean
                for n in nn)
        bestiles = (round((b - tmin) / twidth, 3) * 100 for b in best)
        print(f"{q} by {p}",
              "\t".join(['{:4.3g}']*len(nn)).format(*bestiles),
              sep="\t",file=outfile)

def predict(basis_ym, basis_name, target_ym, target_name, n):
  basis = SS['%s %s'%(basis_ym, basis_name)]
  target = SS['%s %s'%(target_ym, target_name)]
  best = stats.describe(target.all[np.array(basis.iranks_by_all[:n])])
  tbest = stats.describe(target.all[np.array(target.iranks_by_all[:n])])
  m = int((basis.N - n) / 2)
  middle = stats.describe(target.all[np.array(basis.iranks_by_all[m:m+n])])
  tmiddle = stats.describe(target.all[np.array(target.iranks_by_all[m:m+n])])
  randind = list(range(100))
  random.shuffle(randind)
  rlist = np.array(randind[:n])
  randist = stats.describe(target.all[np.array(basis.iranks_by_all)[rlist]])
  trandist = stats.describe(target.all[rlist])
  worst = stats.describe(target.all[np.array(basis.iranks_by_all[-n:])])
  tworst = stats.describe(target.all[np.array(target.iranks_by_all[-n:])])
  print(f"Predicting {basis.id} from {target.id}, window size {n}")
  print("            Prediction        Actual")   
  print("           mean   stdev   mean   stdev  value")
  print("   Best %d %-6.4g %-6.3g  %-6.4g %-6.3g %-6.4g"%(n, best.mean,
                                                     round(math.sqrt(best.variance),
                                                           4),
                                                     tbest.mean,
                                                     round(math.sqrt(tbest.variance),
                                                           4),
                                                     target.all_s.minmax[1]))
  print(" Middle %d %-6.4g %-6.3g  %-6.4g %-6.3g %-6.4g"%(n, middle.mean,
                                                     round(math.sqrt(middle.variance),
                                                           4),
                                                     tmiddle.mean,
                                                     round(math.sqrt(tmiddle.variance),
                                                           4),
                                                     target.all_m))
  print("Randist %d %-6.4g %-6.3g  %-6.4g %-6.3g %-6.4g"%(n, randist.mean,
                                                     round(math.sqrt(randist.variance),
                                                           4),
                                                     trandist.mean,
                                                     round(math.sqrt(trandist.variance),
                                                           4),
                                                     target.all_m))
  print("  Worst %d %-6.4g %-6.3g  %-6.4g %-6.3g %-6.4g"%(n, worst.mean,
                                                     round(math.sqrt(worst.variance),
                                                           4),
                                                     tworst.mean,
                                                     round(math.sqrt(tworst.variance),
                                                           4),
                                                     target.all_s.minmax[0]))

def dump(file=sys.stdout):
  for s in SS.values():
    print("CCStats(%s,%s,%s)"%(repr(s.year_month),repr(s.name),
                               repr(s.file.split('/',7)[7][:-4])),
          file=sys.stdout)

if __name__ == '__main__':
  S = CCStats(*sys.argv[1:])

### I need to review rows, e.g. self.counts[0] is an array of self.N+1 self.counts
###   for the most common label in the complete crawl,
###   from the complete crawl and all the segments
### versus columns, e.g. self.counts[:,0] is an array of self.N decreasing self.counts
###   for all the labels in the complete crawl