#  Time-stamp: <2015-03-01 06:13:00 Ryo KOBAYASHI>
"""
Routines related to neural network.
"""

import math,time,os
import numpy as np
import multiprocessing as mp

#.....import from local modules
import fitpot

#.....constants which may be different from those in fitpot.py
_large= 1.0e+30

#.....global variables
_l1st= True
_nl =  1
_nsp=  1
_nsf=  0
_nhl1= 0
_nhl2= 0
_ergs= []
_frcs= []
_gsf= []
_dgsf=[]
_hl1= []
_hl2= []
_aml= []
_bml= []
_cml= []
_fmatch= False
_basedir='learning_set'
_samples=[]
_sample_dirs=[]
_nprcs= 1
_ergrefs= []
_frcrefs= []
_fmethod= 'test'
_parfile= 'in.params.NN'
_runmode= 'serial'
_rcut  = 5.0
_params= []
_pranges= []

#============================================================ routines
def init(*args,**kwargs):
    """
    Initialize variables for the calculation of NN parameters.
    This should be called at the first place
    before any other NN-related routines.
    """

    global _nl,_nsf,_nhl1,_nhl2,_gsf,_dgsf
    global _fmatch,_basedir,_samples,_sample_dirs,_nprcs
    global _ergrefs,_frcrefs,_fmethod,_parfile,_runmode,_rcut
    global _params,_pranges,_vranges

    _basedir = args[0]
    _params   = args[1]
    _sample_dirs= args[2]
    _samples = args[3]
    _nprcs   = args[4]
    _fmatch  = args[5]
    _ergrefs = args[6]
    _frcrefs = args[7]
    _fmethod = args[8]
    _parfile = args[9]
    _runmode = args[10]
    _rcut    = args[11]
    _pranges = args[12]
    _vranges = args[13]

    #.....read NN parameters
    f= open(_basedir+'/'+'in.const.NN')
    data= f.readline().split()
    _nl  = int(data[0])
    _nsp = int(data[1])
    _nsfc= int(data[2])
    _nhl1= int(data[3])
    if _nl == 2:
        _nhl2= int(data[4])
    n2= 0
    n3= 0
    for line in f.readlines():
        if int(line.split()[0]) == 1:
            n2 += 1
        elif int(line.split()[0]) == 2:
            n3 += 1
    f.close()
    print ' Number of hidden layers             =',_nl
    print ' Number of nodes in hidden-layer-1   =',_nhl1
    if _nl == 2:
        print ' Number of nodes in hidden-layer-2   =',_nhl2
    print ' Number of species                   =',_nsp
    print ' Number of 2-body symmetry functions =',n2
    print ' Number of 3-body symmetry functions =',n3
    
    # check
    ncmb2= _nsp +factorial(_nsp,2)/2
    ncmb3= _nsp *ncmb2
    _nsf= n2*ncmb2 +n3*ncmb3
    print ' Number of symmetry functions        =',_nsf
    if _nl == 1:
        nw= _nsf*_nhl1 +_nhl1
    elif _nl == 2:
        nw= _nsf*_nhl1 +_nhl1*_nhl2 +_nhl2
    if len(_params) != nw:
        print ' [Error] len(params) !=  nw'
        print '   len(params)        = ',len(_params)
        print '   nsf                = ',_nsf
        print '   nhl1               = ',_nhl1