
import os
import socket
import time
import pychemia
from multiprocessing import Process
from pychemia import pcm_log, HAS_PYMONGO

if HAS_PYMONGO:
    from pychemia.db import get_database


class FireballCollector:

    def __init__(self, dbname, source_dir=None, source_file=None, output_names=None, output_file=None,
                 db_settings=None, nconcurrent=1):
        """
        FireballCollector is a class collect Fireball executions into one PyChemia Database.

        The name of the database will be 'dbname' and executions will be processed from a give 'source_dir', a path
        to search for Fireball calculations or 'source_file', one file with one directory per line where Fireball
        executions are stored.
        The variable 'output_names' sets possible output files for Fireball. Typically Fireball output is the standard
        output and the user has to redirect the output to a file with an arbitrary name.

        The processing of directories happens in parallel using 'nconcurrent' python processes.
        By default the database is created on the local machine assuming no authorization or authentication, otherwise
        use the dictionary 'db_settings' to get more control on the database location and access.

        :param dbname:      (str) Name of the database that will contain the Fireball executions
        :param source_dir:  (str) Path for a directory that will be explored recursevely for suitable Fireball
                                  executions
        :param source_file: (str) Path to a file describing directories with Fireball executions, one per line
        :param output_names: (list) List of possible output filenames
        :param output_file: (str) Path to a file with possible output names, one per line
        :param db_settings: (dict) Extra parameters for locating and accessing the mongo server.
        :param nconcurrent: Number of concurrent python executions to fill the database (default: 1)
        """
        self.db_settings = db_settings
        self.dbname = dbname
        self.source_dir = source_dir
        self.source_file = source_file
        self.output_names = output_names
        self.output_file = output_file
        self.nconcurrent = nconcurrent
        self.sleeping_time = 2
        self.database = None

        if self.output_file is not None and os.path.exists(self.output_file):
            rf = open(self.output_file)
            self.output_names = [x.strip() for x in rf.readlines()]

        if self.source_dir is None and self.source_file is None:
            raise ValueError("One of the variables source_dir or source_file need to be defined")

        if self.db_settings is None:
            self.db_settings = {'name': self.dbname}
        else:
            self.db_settings['name'] = self.dbname

        self.db_settings = dict(self.db_settings)
        # The first component of each pair in to_evaluate is the name of the database
        self.db_settings['name'] = self.dbname

        self.database = get_database(self.db_settings)

    def worker(self, path):
        """
        From a given 'path', the worker will collect information from several files
        such as:

        fireball.in
        eigen.dat
        param.dat
        answer.bas
        output collected from the standard output of fireball.x


        After collecting data into dictionaries, the data is stored in a PyChemia database
        """

        pcdb = get_database(self.db_settings)
        files = os.listdir(path)
        properties = {}
        geo = None

        if os.path.lexists(path+os.sep+'fireball.in'):
            properties['path'] = path
            score = 'input'
            try:
                invars = pychemia.code.fireball.read_fireball_in(path + os.sep + 'fireball.in')
                properties['input'] = invars
            except:
                print('Bad fireball.in on %s' % path)

            if os.path.exists(path+os.sep+'param.dat'):
                try:
                    score += '+param'
                    param = pychemia.code.fireball.read_param(path + os.sep + 'param.dat')
                    properties['param_dat'] = param
                except:
                    print('Bad param.dat')

            if os.path.lexists(path+os.sep+'eigen.dat'):
                try:
                    score += '+eigen'
                    eigen = pychemia.code.fireball.read_eigen(path + os.sep + 'eigen.dat')
                    properties['eigen_dat'] = eigen
                except:
                    print('bad eigen.dat')

            if os.path.lexists(path+os.sep+'answer.bas'):
                try:
                    score += '+answer'
                    geo = pychemia.code.fireball.read_geometry_bas(path + os.sep + 'answer.bas')
                    periodic = False
                except:
                    print('Bad answer.bas')

            for ioutput in self.output_names:
                if os.path.isfile(path + os.sep + ioutput):
                    rf = open(path + os.sep + ioutput)
                    line = rf.readline()
                    rf.close()
                    if "Welcome to FIREBALL" in line:
                        try:
                            output = pychemia.code.fireball.read_final_fireball_relax(path + os.sep + ioutput)