#!/usr/bin/env python
import os

import click

@click.group()
def cli():
    pass
    
@click.command(name='check_input')
@click.option('-i', '--input_path', required=True, help='path to text file containing names of input samples')
def check_input(input_path):
    tool = SampleTool()
    samples = tool.check_mc_samples(input_path)
    daod_samples = tool.check_DAOD_from_AOD(input_path)
    print("----------------------------------------- DAOD INPUT CHECK -----------------------------------------")
    print("INFO: Avaliable Samples:")
    print("\n".join(samples['exist']))
    print("")
    print("INFO: Missing Samples:")
    print("\n".join(samples['missing']))
    print("")
    print("INFO: The following samples have corresponding derived samples:")
    for AOD, DAOD in daod_samples.items():
        print("{}:\n{}".format(AOD, "\n".join(["\t{}".format(d) for d in DAOD])))
    print("----------------------------------------------------------------------------------------------------")
    
    
@click.command(name='list_samples')
@click.option('-n', '--name', 'pattern', required=True, show_default=True,
              help='Expression used to filter the DAODs. Wild card is accepted.')
@click.option('-t', '--did-type', default='container', show_default=True,
              help='Filter by DID type.')
@click.option('--not-empty/--allow-empty', default=True, show_default=True,
              help='Whether to filter empty DAODs.')
@click.option('--latest-ptag/--any-ptag', default=True, show_default=True,
              help='Wheter to show only the DAODs with the latest ptags.')
@click.option('--single-rtag/--allow-multi-rtag', default=True, show_default=True,
              help='Wheter to show only the DAODs with a single rtag.')
@click.option('--single-ptag/--allow-multi-ptag', default=True, show_default=True,
              help='Wheter to show only the DAODs with a single ptag.')
@click.option('--esrp-tags-only/--allow-any-tags', default=True, show_default=True,
              help='Wheter to show only the DAODs with e, s, r and p tags only.')
@click.option('--data-type', default='', show_default=True,
              type=click.Choice(["DAOD", "NTUP_PILEUP", ""], case_sensitive=False),
              help='Filter samples by data type. Choose from "DAOD", "NTUP_PILEUP", "".')
@click.option('--display-columns', default="did,nevent,size", show_default=True,
              help='Columns to display separated by commas. Common columns include "did", "did_type"'
              ', "ptag", "nevent", "size", etc.')
@click.option('--show-table/--show-text', default=True, show_default=True,
              help='Whether to display table or just plain text of dids.')
@click.option('--outname', '-o', default=None, show_default=True,
              help='If specified, save the output with the given filename.')
def list_samples(**kwargs):
    aux_kwargs = {}
    for key in ['display_columns', 'outname', 'show_table']:
        aux_kwargs[key] = kwargs.pop(key)
    from derivation_tool import DIDTool, DIDCollection
    tool = DIDTool()
    df = tool.list_derived_samples(**kwargs)
    collection = DIDCollection(df)
    display_columns = [i.strip() for i in aux_kwargs['display_columns'].split(',')]
    if ('nevent' in display_columns) or ('size' in display_columns):
        collection.fill_attributes(file_data=True, inplace=True)
    if aux_kwargs['show_table']:
        collection.print_table(attributes=display_columns)
    else:
        for did in collection.get_dids():
            print(did)
    outname = aux_kwargs['outname']
    if outname is not None:
        export_format = os.path.splitext(outname)[1].strip(".")
        if export_format == 'csv':
            collection.df.to_csv(outname, index=False)
        elif export_format == 'json':
            data = collection.df.to_dict('records')
            import json
            json.dump(data, open(outname, 'w'), indent=2)
        else:
            dids = collection.get_dids()
            with open(outname, 'w') as f:
                for did in dids:
                    f.write(f"{did}\n")

if __name__ == "__main__":
    cli.add_command(list_samples)
    cli()   