#!/usr/bin/python
"""
This script reads n-resolved cross section data, as generated by "Stueckelberg",
and performs the l-distribution low-energy and statistical approximations on
that data.

usage: /home/micha/bin/ldist.py [-h] [-l] [-s] [-v] file

General l-distribution estimations.

positional arguments:
  file        Input file.

optional arguments:
  -h, --help  Show help message and exit.
  -l          Low-energy l-distribution.
  -s          Statistical l-distribution.
  -v          Verbose. Intended for debugging.

$Id: ldist.py,v 1.2 2019/03/09 20:01:28 micha Exp micha $
$Author: micha $
$Revision: 1.2 $
$Log: ldist.py,v $
Revision 1.2  2019/03/09 20:01:28  micha
Corrected mistake in low-energy factor.

"""
import numpy as np
import os
import argparse
import sys

class LDist(object):
    def __init__(self,fname):
        self.fname=fname
        self.basename=os.path.splitext(self.fname)[0]
        self.data=self.GetData()
        self.header=self.GetHeader()
        self.labels=self.GetLabels()
        self.nlist=self.Getn()
        self.r,self.c=self.data.shape
        return
    def GetData(self):
        return np.loadtxt(self.fname)
    def GetHeader(self):
        header=[]
        with open(self.fname,"r") as f:
            for l in f:
                if(l.strip().startswith("#")):
                    header.append(l)
        f.close()
        return header
    def GetLabels(self):
        for l in self.header:
            if(l.startswith("# (eV/u)")):
                labels=l.split()
        return labels[2:-1]
    def Getn(self):
        import re
        return [int(x[0]) for x in [re.findall(r'\d+',i) for i in self.labels] if x]
    def cf(self,n,l):
        llist=['s','p','d','f','g','h','i','k','l','m','n','o','q','r','t','u','v','w','x','y','z']
        return str(n)+llist[l]
    def Wst(self,n,l):
        return (2*l+1.0)/n**2
    def Wle(self,n,l):
        from math import factorial as fact
        return ((2*l+1.0)*(fact(n-1))**2)/(fact(n+l)*fact(n-l-1))
    def ldist(self,W):
        newdata=np.vstack(self.data[:,0])
        fmtstr="%9.3e"
        labstr="# (eV/u)"
        for i,n in enumerate(self.nlist):
            for l in reversed(range(n)):
                if(W=="stat"):
                    newdata=np.concatenate((newdata,np.vstack(self.Wst(n,l)*self.data[:,i+1])),1)
                elif(W=="lowe"):
                    newdata=np.concatenate((newdata,np.vstack(self.Wle(n,l)*self.data[:,i+1])),1)
                fmtstr+="\t%13.7e"
                labstr+="\t"+self.cf(n,l)+"\t"
        newdata=np.concatenate((newdata,np.vstack(np.sum(newdata[:,1:],axis=1))),1)
        fmtstr+="\t%13.7e"
        labstr+="\tTotal\n"
        if(W=="stat"):
            self.stdata=newdata
            ofile=self.basename+"_stat.cs"
        elif(W=="lowe"):
            self.ledata=newdata
            ofile=self.basename+"_lowe.cs"
        f=open(ofile,"w")
        for h in self.header[:-1]:
            f.write(h)
        f.write(labstr)
        f.close()
        f=open(ofile,"a")
        np.savetxt(f,newdata,delimiter="\t",fmt=fmtstr)
        f.close()
        return

def MainMenu(argv):
    parser=argparse.ArgumentParser(description="l-distribution estimations.",prog=sys.argv[0])
    parser.add_argument("InFile",help="Input file.",type=str,metavar="file",default="/dev/stdin")
    parser.add_argument("-l",help="Low-energy l-distribution.",action="store_true",dest="lowe",default=False)
    parser.add_argument("-s",help="Statistical l-distribution.",action="store_true",dest="stat",default=False)
    parser.add_argument("-v",help="Verbose. Intended for debugging.",action="store_true",dest="Verbose",default=False)
    parser._actions[0].help="Show help message and exit."
    args=parser.parse_args(None if argv else ["-h"])
    if(args.Verbose):
        for i in vars(args):
            print("%s=%s"%(i,getattr(args,i)))
    return args

def main(argv):
    a=MainMenu(argv)            # Display the main menu and get args.
    ld=LDist(a.InFile)          # Create an instance (object) of LDist;
                                # a lot of processing is being done here.
    if(a.stat):ld.ldist("stat") # Perform the Statistical l-distribution.
    if(a.lowe):ld.ldist("lowe") # Perform the Low-energy l-distribution.

if __name__=="__main__":
    main(sys.argv[1:])
