#!/usr/bin/env python

import itertools
import types
import ROOT
from array import array
from math import log10
from . import logger

#--------#
# SCHEME #
#--------#

def COLORS(i):
  """
  Repeating palette of colors, modulo around.
  Change to obey the active palette, allowing local change of color.

  >>> COLORS(0)
  1
  >>> COLORS(3)  # qhist's
  28

  """
  n = ROOT.TColor.GetNumberOfColors()
  return ROOT.TColor.GetColorPalette(i%n)
  # return (1, 2, 4, 28, 6, ROOT.kGreen+3, ROOT.kAzure+7 )[i%7]
  # return ( 1,2,8,4,6,7,9 )[i%7]
  # return ROOT.gStyle.GetColorPalette( i*100 )
  # return (0,5,7,3,6,2,4,1)[i%8] ## From LHCb.style

def MARKERS(i):
  """
  Repeating choices of marker.

  >>> MARKERS(0) == ROOT.kMultiply
  True
  >>> MARKERS(2) == ROOT.kPlus
  True

  """
  return [ROOT.kMultiply, ROOT.kCircle, ROOT.kPlus, ROOT.kOpenDiamond, ROOT.kOpenTriangleUp][i%5]

def LINES(i):
  """
  Choice of line style.

  >>> LINES(0)
  1

  """
  return i+1

def FILLS(i):
  """
  Repeating choice of fill style, modulo 7.
  
  >>> FILLS(0)
  3001

  """
  return (3001, 3254, 3245, 3205, 3295, 3012, 3016)[i%7]
  # return (3001, 3254, 3245, 3305, 3395, 3012, 3016)[i%7]
  # return (3003,3004,3005,3006,3007,3012,3016)[i%7]

#-------------#
# COMPUTATION #
#-------------#

def auto_bin_size(nevt):
  """
  Given the list of event size (counts), return appropriate num of bins,
  which is calculated from sqrt of smallest event size, rounded in 5.
  min of 20, max of 100.

  >>> auto_bin_size(100)
  20
  >>> auto_bin_size(5000)
  70
  >>> auto_bin_size(1E6)
  100
  >>> auto_bin_size(0)  # Even null plot, return minimum number of bins
  20

  """
  minv = 20
  maxv = 100
  val  = int(round(1./5*(nevt**0.5))*5)
  return minv if val < minv else maxv if val > maxv else val

#-------------------------------------------------------------------------------

def auto_log(vmin, vmax):
  """
  Deduce whether the log scale should be used or not, judging from given range.

  >>> auto_log(0, 10)  # has zero, don't.
  False
  >>> auto_log(0.0, 1.0)  # has zero, don't.
  False
  >>> auto_log(-1.0, 1.0)  # has negative, don't.
  False
  >>> auto_log(1E-4, 1E1)  # positive & wide, do it.
  True
  >>> auto_log(10, 50)  # arbitary, don't
  False

  >>> ## If both are not initialize (e.g., null-stats), return None 
  >>> ## in order to have same default as QHist.xlog
  >>> auto_log( 10, None ) is None
  True
  >>> auto_log( None, 100 ) is None
  True
  >>> auto_log( None, None ) is None
  True

  """
  ## Uninitialize case
  if vmin is None or vmax is None:
    return None
  ## If bound within [0,1], don't
  if vmin==0. and vmax==1.:
    return False 
  ## Crossing zero. No log
  if vmin <= 0:
    return False 
  ## Large magnitude
  if vmax / vmin >= 1E2:
    return True 
  ## Don't do log, by default
  return False

#-------------------------------------------------------------------------------

def auto_morelog(axis):
  """
  Wrapper around TAxis.SetMoreLogLabels 
  Configure automatically whether this axis should have detail label or not.
  Perferable do so in small scale.

  >>> axis = ROOT.TAxis(20, 1, 1000)
  >>> axis.moreLogLabels
  False
  >>> auto_morelog(axis)
  >>> axis.moreLogLabels
  True

  """
  valmin = axis.xmin
  valmax = axis.xmax
  if valmin and valmax/valmin < 1E4: # Non zero, and not-too-wide scale
    axis.moreLogLabels = True

#-------------------------------------------------------------------------------

def make_bins(valmin, valmax, nbins, is_log=False, margin=0.0):
  """
  Important method to provide the binnings for given spec.
  The cue is to have pretty-log axis.
  http://root.cern.ch/root/roottalk/roottalk06/1213.html

  Usage::

      ## Simple
      >>> make_bins(0, 100, 20)
      array('d', [0.0, 5.0, 10.0, ...

      ## Restrict the extreme one
      >>> make_bins(-1e12, 1, nbins=1)
      array('d', [-1000000.0, 1.0])
      >>> make_bins(1, 1e12, nbins=1)
      array('d', [1.0, 1000000.0])

      ## Grace allowance
      >>> make_bins(10, 10, nbins=1)
      array('d', [9.9, 10.1])
      >>> make_bins(0, 0, nbins=1)
      array('d', [-0.1, 0.1])

      ## Ignore log on negative
      >>> make_bins(-1, 10, nbins=1, is_log=True)
      array('d', [-1.0, 10.0])

      ## absurd
      >>> make_bins(-1e12, -1e12, 100)
      Traceback (most recent call last):
      ...
      ValueError: Data range seems absurd. The script will halt.

  """
  ### Validate and abort if it seems very suspicious
  if valmin < -1E6 and valmax < -1E6:
    logger.warning(locals())
    raise ValueError('Data range seems absurd. The script will halt.')

  ### Some fine adjustment of data
  ## If it's too extreme, restrict the domain
  if valmin < -1E6:
    logger.debug('valmin too small. Restrict back to -1E6.')
    valmin = -1E6
  if valmax > 1E6:
    logger.debug('valmax too large. Restrict back to 1E6.')
    valmax = 1E6

  ## If the min==max, do the grace allowance (+-1% each side)
  if valmin == valmax:
    if valmin==0:  # Worst case: Complete-0s
      dval = 0.1
    else:
      dval = abs(valmin*0.01)
    valmin -= dval
    valmax += dval
    logger.debug('Grace adjust when min==max: dx={}'.format(dval))

  ## If x is logarithm, disable if math domain conflict
  if is_log and valmin <= 0:
    logger.warning('Found non-positive valmin in log-axis. Recommending disable log-axis, or put `xmin=1E-6` or `cuts>0`)')
    is_log = False

  logger.debug('## QHistUtils.make_bins summary (post-adjustment)')
  logger.debug("valmin : " + str(valmin))
  logger.debug("valmax : " + str(valmax))
  logger.debug("nbins  : " + str(nbins))
  logger.debug("is_log : " + str(is_log))

  ## Finally, make a array out of it.
  valmin *= ( 1.-margin if valmin>0 else 1.+margin )
  valmax *= ( 1.+margin if valmax>0 else 1.-margin )
  if is_log:
    width = 1.0*(log10(valmax)-log10(valmin))/(nbins)
    arr   = [ 10**(log10(valmin) + i*width) for i in range(nbins+1) ]
  else:
    width = 1.0*(valmax-valmin)/(nbins)
    arr   = [ valmin + i*width for i in range(nbins+1) ]
  return array('d', arr)

#-------------------------------------------------------------------------------

def get_stats_single(t, p, c):
  """
  Return the axis-min/max of the given TPC tuple.

  Note: TTree::GetMinimum is NOT what we want because cannot use cut c
  Note: Y-min/max is more complicate because of ybin, normalization
  Note: For SetEstimate hack, see TTree::Draw documantation
  Note: To ensure, use the vanilla TTree (not rootpy Tree), force mother call.

  """
  logger.debug('Stats: t: '+t.title)
  logger.debug('Stats: p: %r'%p)
  logger.debug('Stats: c: %r'%c)
  t.estimate = t.entries+1 # bug with (-1)
  # n = ROOT.TTree.Draw(t, p, c, "goff")
  # n = t.Draw(p, c, "goff")
  cw = c.replace('&&','*').replace('&','*')
  n  = t.Draw(p, cw, "goff") # make sure it counts the weight
  if n <=0 :
    logger.warning("Null-histogram. Please know what you're doing...")
    logger.warning('Tree : %s %r'%(t.title, t))
    logger.warning('Param: %s'%p)
    logger.warning('Cut  : %s'%c)
    return 0, None, None
    # raise ValueError('Null plots!! Please recheck the TPC settings.')
  g1,g2 = itertools.tee(t.GetV1()[i] for i in xrange(n)) # Use generator for speed
  return n, min(g1), max(g2)
  # axis = ROOT.gROOT.FindObject('htemp').GetXaxis()
  # return n, axis.GetXmin(), axis.GetXmax()

#-------------------------------------------------------------------------------

def prefer_first(*vals):
  """
  This idiom is used so often...
  Pick fist value unles it's None.

  Usage::
  
      >>> prefer_first( True, 42 )
      True
      >>> prefer_first( None, 42 )
      42

      ## Still pick 1st one if it's not None
      >>> prefer_first( False, 42 )
      False

      >>> prefer_first( 1, 2, 3 )
      1
      >>> prefer_first( None, 2, 3 )
      2
      >>> prefer_first( None, None, 3 )
      3
      >>> prefer_first( None, None, None ) is None
      True

  """
  for val in vals:
    if val is not None:
      return val 
  return val # last value in loop

#--------#
# STRING #
#--------#

def cut_minmax(p, valmin, valmax):
  """
  Simply parse 3 strings together in form of 'valmin < p && p < valmax'.
  Skip if val is None. Return empty string if both null.

  >>> cut_minmax( 'X', 10, 20 )
  '(X >= 10) & (X <= 20)'

  >>> cut_minmax( 'X', 12.3, None )
  'X >= 12.3'

  >>> cut_minmax( 'X', None, None )
  ''

  """
  queue = []
  if valmin is not None: queue.append(p+' >= '+str(valmin))
  if valmax is not None: queue.append(p+' <= '+str(valmax))
  return join(queue)

#-------------------------------------------------------------------------------

def pretty_array(arr):
  """
  Pretty-print an array('d') to not be too long, for logging.

  >>> from array import array
  >>> arr = array('d', [10*x**-2 for x in xrange(5,10)])
  >>> arr
  array('d', [0.4, 0.2777777777777778, 0.2040816326530612, 0.15625, 0.12345679012345678])
  >>> print(pretty_array(arr))
  array(0.40, 0.28, 0.20, 0.16, 0.12)

  """
  l = [ '%.2f'%d for d in arr ]
  s = str(l).replace("'",'').replace('[','').replace(']','')
  s = s[:75] + (s[75:] and '..')
  return "array(%s)" % s

#-------------------------------------------------------------------------------

def pretty_round(val):
  """
  Return string used in mean/RMS, with my custom format.

  >>> pretty_round(None)      # Nice null
  '---'
  >>> pretty_round(1234567)   # Too large
  '> 1E6'
  >>> pretty_round(-1234567)  # Too large
  '< -1E6'
  >>> pretty_round(2E-7)      # Tiny than 1E-6
  '< 1E-6'
  >>> pretty_round(3E-5)      # Small
  '3.00E-05'
  >>> pretty_round(12345)     # Large 
  '1.23E+04'
  >>> pretty_round(1.4142)    # Natural scale
  '1.41'

  """
  if val is None:
    return '---'
  if val > 1E6:  # So largely positive
    return '> 1E6'
  if val < -1E6: # So largely negative
    return '< -1E6'
  if 0 < val < 1E-6:  # So tiny near zero
    return '< 1E-6'
  if 0 < val < 1E-2:  # semi-small value
    return '{:.2E}'.format(val)
  if 1E4 < val < 1E9:  # semi-large value
    return '{:.2E}'.format(val)
  return "{:.2f}".format(val) # Human-friendly format

#-------------------------------------------------------------------------------

def concat(*tokens, **kwargs):
  """
  Custom method to concatenate list of string together, such that
  null is ignore ( being basestring or not ).
  Provide flag `delim` to change delimiter

  Use primarily on logic of QHist's naming

  >>> concat('prefix', 'root', 'suffix')
  'prefix_root_suffix'

  >>> concat(None, 'root', 'suffix', delim='/')
  'root/suffix'

  """
  delim = kwargs.get('delim', '_')
  return delim.join([ str(token).strip(delim) for token in tokens if token ])

#-------------------------------------------------------------------------------

def join(*cuts):
  """
  Given list of string (cuts), join them and parenthesis guard 
  (guard only the inner element, not the full outer, to avoid verbose double-guard)
  If the iterable is LIST/TUPLE/GEN, join with AND 
  If the iterable is SET           , join with OR 
  Designed for recursive usage.

  Note:
  - No using double-op '&&' to be LoKi-compatible.
  - The root argument, *args, by python's design, is intrinsically a tuple.
    So that's why one can pass arguments instead of list-of-string as arg.

  Args:
    cuts (iterable): Iterable (tuple/list/set) of string to join.

  Usage::

      ## Test basics
      >>> join([ 'a', 'b' ]) == join('a','b')
      True
      >>> join( 'a', 'b' )
      '(a) & (b)'
      >>> join({'a', 'b'})
      '(a) | (b)'

      ## Handle effective null
      >>> join()
      ''
      >>> join([])
      ''
      >>> join([''])
      ''
      >>> join([None])
      ''
      >>> join({''})  # List of null-cuts
      ''
      >>> join([('',),{''}])  # Super null!
      ''
      >>> join(['c > 0', ('',), ''])
      'c > 0'

      ## Simplify idempotent nesting
      >>> join([[( 'a > 3' )]])
      'a > 3'
      >>> join([ {'b1','b2'} ])
      '(b1) | (b2)'

      ## Handle nested
      >>> join( 'a', {'b1','b2'}, 'c' )
      '(a) & ((b1) | (b2)) & (c)'

      ## Trim whitespace 
      >>> join( 'a     > 5', 'b <    10' )
      '(a > 5) & (b < 10)'

  """

  ## Early kill the null-value on any case, string or not:
  if not cuts:
    return ''
  ## De-equip if it's single-tuple
  if isinstance( cuts, tuple ) and len(cuts)==1:
    cuts = cuts[0]
  ## Early deal with recursive: String is the atomic unit 
  if isinstance( cuts, basestring ):
    ## Simple aesthetic on string: Trim all multiple-whitespace to single one
    return ' '.join(cuts.split())

  ## Given now an iterable, determine the operator 
  #  (need this info before the cleansing).
  if isinstance( cuts, (list,tuple,types.GeneratorType) ):
    op = ' & '
  elif isinstance( cuts, set ):
    op = ' | '
  else: # pragma: no cover
    raise ValueError('Invalid iterable type: %r' % type(cuts))

  ## Apply recursion here, kill bad one before passing.
  cuts = [ join(c) for c in cuts if c ]
  ## Kill bad ones after recursion is applied
  cuts = [ c for c in cuts if c ]
  ## Empty list is discarded at this point
  if not cuts:  
    return ''
  ## Early deal with collection with one item: Extract it.
  if len(cuts) == 1:
    return cuts[0]
  ## Finally join with current op. Wrap inner item here, and only here
  return op.join([ '(%s)'%c for c in cuts])

#-------------------------------------------------------------------------------

_SAFENAME_EXTENSIONS = 'pdf', 'png', 'eps', 'c'

def safename(s):
  """
  Return safename for ROOT.TFile creation

  >>> safename('Hello(world)')  # ROOT use bracket for something else
  'Hello{world}'

  >>> safename('PT/P')  # Bad for sys
  'PT_over_P'

  >>> safename('TMath::Log10(PT)')
  'Log10{PT}'

  >>> safename('P/1e3')
  'P'
  """
  ## remove const
  blacklist = ' ', '/1000', '/1e3', '/1E3', '*1000', '*1e3', '*1E3', '*1e6', '*1E6'
  for x in blacklist:
    s = s.replace(x, '')
  ## replace bad ones
  s = s.replace('(', '{').replace(')','}')  # Danger for root
  s = s.replace('/','_over_')
  s = s.replace('TMath::', '')  # Nicer name
  s = s.replace(':', '__') # APT:M --> APT__M
  s = s.replace('.', '_')
  for ex in _SAFENAME_EXTENSIONS:
    s = s.replace('_'+ex.upper(),'.'+ex.upper())
    s = s.replace('_'+ex,'.'+ex)
  return s

#-------------------------------------------------------------------------------

def shorten(s, width=78):
  """
  Helper method to shorten string (to 78 chars, terminal width) 
  for pretty-printing.

  >>> print(shorten('abcdefghijklmnopqrstuvwxyz'))  # not 78
  abcdefghijklmnopqrstuvwxyz
  
  >>> print(shorten('abcdefghijklmnopqrstuvwxyz'*10))
  abcdefghijklmnopqrstuvwxyzabcdefghijkl...pqrstuvwxyzabcdefghijklmnopqrstuvwxyz
  
  """
  s = str(s)
  return s if len(s)<width else (s[0:width/2-1]+'...'+s[len(s)-width/2+2:])

#-------------------------------------------------------------------------------

def split_dim(p):
  """
  Safe method to split multi-dim params string.
  Do not change any order

  >>> split_dim('MM:PT')
  ['MM', 'PT']
  >>> split_dim('TMath::Log10(MM): PT')
  ['TMath::Log10(MM)', 'PT']

  """
  l = p.replace("::","@").split(":")
  return [ x.replace("@","::").strip() for x in l ]

#-------------------------------------------------------------------------------

def dim(p):
  """
  Short-cut method of above: Return integer dimension of given string.
  Note that in this case, it can detect dim-0

  >>> dim('PT')
  1
  >>> dim('PT:PZ')
  2
  >>> dim('')
  0
  >>> dim( 'TMath::Log10(PT)' )
  1
  >>> dim( 'TMath::Log10(PT):PZ' )
  2
  """
  return len(split_dim(p)) if p else 0

#===============================================================================
