#!/usr/bin/env python 

""" 
Analyse the disk space usage of an aanet file.

 usage analyse_file_usage.py -f input_file

options:
 -f : input file        default=/sps/km3net/users/coniglio/antares_seawiet/oct17/MC_v5/mc/atm_neutrino/KM3NeT_-00000001_20171212/v5.1_res/reco/mcv5.1_res.genhen_numuCC.km3_AAv1.jte.jchain.aashower.1.root
 -g : graphical summary output                       default=True
 -n : normalize histogram                            default=False
 -h : this help message and exit                     default=False       
 -i : Python interative mode (prompt when done)      default=True  
"""


import sys,aa; 
from ROOT import TFile
import sys
options = aa.Options( __doc__, sys.argv[1:] ) 


def autotype( v ):

    "why oh why is this not in the standard lib?"
    for t in (int, float, str) : 
        try : return t(v)
        except ValueError: pass
    return v

class Redirector( object ):
    "from https://root.cern.ch/phpBB3/viewtopic.php?f=14&t=8605"
    def __init__( self, stream ):
        import os, sys, tempfile
        w, n = tempfile.mkstemp()
        os.unlink( n )
        self._file = os.fdopen( w, 'r+' )
        self._stream = stream

        self._savefn = os.dup( self._stream.fileno() )
        os.dup2( self._file.fileno(), self._stream.fileno() )

    def restore( self ):
        import os
        
        self._stream.flush()
        os.dup2( self._savefn, self._stream.fileno() )
        
        self._file.seek( 0 )
        content = ''.join( self._file.readlines() )
        self._file.close()
        del self._file, self._savefn, self._stream
        
        return content



rootfile = TFile (options.f)
tree = rootfile.Get("E")


r = Redirector( sys.stdout )
tree.Print()
s = r.restore()


L , record = [], {}

def getnext( line, word , skip = 1) :

    v = line.split();

    try : 
        return v[ v.index(word)+skip];
    except ValueError : 
        return None




from copy import copy

for x in s.splitlines() :
    
    if x.startswith("*....") :
        if 'branchname' in record.keys() :
            L.append( copy(record) )
        record = {}

    def sett( name , key , skip =1 ) :

        y = getnext( x, key, skip )
        if y : record[name] = y


    x = x.replace("*","").replace(":","").replace("=","")
    print (x)

    sett( "entries",    "Entries",1 )
    sett( "branchname", "Br",     2 )
    sett( "memsize",    "Total",  2 )
    sett ("filesize",   "File",   2 )

    if not 'branchname' in record.keys() : continue

    v = record['branchname'].split(".")
    if len(v) > 1 and v[0] != 't' : 
        record['collection'] = v[0]
    else :
        record['collection'] = "evt"




vars = "branchname collection entries memsize filesize".split()
typs = "C          C          I       I       I".split()

descriptor = ":".join( v[0]+"/"+v[1] for v in zip(vars, typs) )

import ROOT
t = ROOT.TTree("t","t")


s = ""
for r in L :

    print (r)

    if not 'branchname' in r.keys() : continue
    if not 'filesize' in r.keys() : r['filesize'] = "0"
 #   try :   
    s += " ".join( r[v]  for v in vars ) +"\n"
 #   except KeyError : 
 #       pass

ss = ROOT.istringstream(s)
t.ReadStream(ss, descriptor)

if options.g :
    t.SetFillColor(5)
    t.Draw("collection>>hh","filesize","goff")
    ROOT.gStyle.SetOptStat(0)
    if options.n : 
        ROOT.hh.Scale( 1.0 / ROOT.hh.Integral() )
        ROOT.hh.SetMinimum(1e-4)
    else :
        ROOT.hh.SetMinimum(1000)
    ROOT.hh.SetTitle( options.f.split("/")[-1] )
    ROOT.hh.SetYTitle("fraction of disk space")
    
    ROOT.hh.Draw("hist")
    ROOT.gPad.SetLogy()
    ROOT.gPad.Update()