#!/usr/bin/env python
## ...........................................................................
## 
## PIMMS (Polymer Interactions in Multicomponent Mixtures)
## Alex Holehouse, Pappu Lab, Holehouse Lab 
## Copyright 2015 - 2020
## ...........................................................................

import sys
import argparse 
import os

from pimms.simulation import Simulation
from pimms.keyfile_parser import KeyFileParser
from pimms import IO_utils, CONFIG


import pimms 

def print_startup():
    IO_utils.status_message('','major')
    print("+-------------------------------------------------+")
    print("|                                                 |")
    print("|                      PIMMS                      |")
    print("|                                                 |")
    print("| Polymer Interactions in Multicomponent Mixtures |")
    print("|                                                 |")
    print("+-------------------------------------------------+")
    print("         Version [%s]" % (pimms.__version__))
    print("")
    IO_utils.status_message('STARTING PIMMS...','major')

def print_info_header():
    print_startup()

    IO_utils.stdout('[-i/--info] flag prints out information on every keyword that a PIMMS simulation can use. Flags are divided into those which are required for PIMMS to run, and those which modulate default behaviour. \n', 60)
    #IO_utils.stdout('This is not yet complete - sorry!\n', 60)


def grouped_keywords():
    """Return [(heading, [keywords]), ...] for `--info`, grouped per
    CONFIG.KEYWORD_GROUPS, with any documented keyword not assigned to a group
    collected under a trailing 'Other' heading so nothing is ever dropped."""
    groups = []
    seen = set()
    for heading, keys in CONFIG.KEYWORD_GROUPS:
        present = [k for k in keys if k in CONFIG.KEYWORDS_DESCRIPTION]
        if present:
            groups.append((heading, present))
            seen.update(present)
    leftover = [k for k in CONFIG.KEYWORDS_DESCRIPTION if k not in seen]
    if leftover:
        groups.append(("Other", leftover))
    return groups


# because we want this to be a stand alone program we need to include all
# the argparse stuff inside a __main__ guard
if __name__=="__main__":

    # We create a argument parser (called 'parser') here 
    # Basically what this does is create an ArgumentParser object
    parser = argparse.ArgumentParser()

    # We add the command line argument(s) we expected to parse
    # I've only defined one here, but we can add as many as we want.
    #
    #
    # --file is the long name which the program uses
    # -f is the short name
    # help describes what this argument does

    parser.add_argument("-keyfile", "-k", help="Filename of keyfile") 
    parser.add_argument("--info", "-i", help="Prints a lot of keyfile options", nargs='?', const='LIST') 
    parser.add_argument("--version", "-v", help="Prints out the current version", action='store_true') 

    # having defined our argument variables we now parse the arguments parsed
    # to the program
    args = parser.parse_args()


    # if args has a variable called keyfile (which will happen if we passed
    # our program an argument associated with the -k or -keyfile flag
    if args.version:
        print('version %s' % (pimms.__version__))
        
    elif args.info is not None:
        print_info_header()
        
        multiline_leader='          '
        
        if args.info in CONFIG.KEYWORDS_DESCRIPTION:
            k = args.info
            k_type    = CONFIG.KEYWORDS_DESCRIPTION[k][0]
            k_desc    = CONFIG.KEYWORDS_DESCRIPTION[k][1]

            try:
                k_default = CONFIG.DEFAULTS[k]
            except KeyError:
                k_default = 'N/A'

                

            IO_utils.stdout(f'NAME    : {k}', multiline_leader=multiline_leader)
            IO_utils.stdout(f'TYPE    : {k_type}', multiline_leader=multiline_leader)
            IO_utils.stdout(f'DESCR   : {k_desc}', multiline_leader=multiline_leader)
            IO_utils.stdout(f'DEFAULT : {k_default}', multiline_leader=multiline_leader)
            IO_utils.newline()
        elif args.info == 'ALL':
            for heading, present in grouped_keywords():
                IO_utils.newline()
                print('===== %s =====' % heading)
                IO_utils.newline()
                for k in present:
                    k_type = CONFIG.KEYWORDS_DESCRIPTION[k][0]
                    k_desc = CONFIG.KEYWORDS_DESCRIPTION[k][1]

                    try:
                        k_default = CONFIG.DEFAULTS[k]
                    except KeyError:
                        k_default = 'N/A'

                    IO_utils.stdout('NAME:     %s' % (k), multiline_leader=multiline_leader)
                    IO_utils.stdout('TYPE:     %s' % (k_type), multiline_leader=multiline_leader)
                    IO_utils.stdout('DESC:     %s' % (k_desc), multiline_leader=multiline_leader)
                    IO_utils.stdout('DEFAULT : %s' % (k_default), multiline_leader=multiline_leader)
                    IO_utils.newline()
        elif args.info == 'LIST':
            IO_utils.stdout('Keywords are grouped by purpose below. For details on one pass -i <keyword name>, or pass -i ALL to print every keyword description.')
            print('')
            for heading, present in grouped_keywords():
                print(heading)
                print('-' * len(heading))
                IO_utils.stdout('  '.join(present), multiline_leader='  ')
                print('')
            
                
            
        else:
            print(f"Unknown keyword '{args.info}'. Run 'PIMMS --info' to list every valid keyword, "
                  f"'PIMMS --info <keyword>' for details on one, or 'PIMMS --info ALL' to print them all.")
    
        
    elif args.keyfile:

        # print the standard startup message!
        print_startup()

        if not os.path.isfile(args.keyfile):
            IO_utils.status_message('Could not open file "%s". Exiting...\n' %(args.keyfile), msg_type='error')
            exit(1)
            
        # read in the keyfile
        KeyFile = KeyFileParser(args.keyfile)

        # print a summary
        KeyFile.print_summary()
        
        # Set up a new simulation
        Simulator = Simulation(KeyFile.keyword_lookup)

        # Run a new simulation
        Simulator.run_simulation()

    else:
        print('Pass all arguments (try --help for more info)')


