import os
from typing import Optional, List, Dict

from hpogrid.components import HPOReport

class PandaTaskReport(HPOReport):
    
    def __init__(self):
        super().__init__()
        from hpogrid.components.panda_task_manager import PandaTaskManager
        self.taskmgr = PandaTaskManager()
        
    def display(self, project_name, outname="hpo_result", metric_min=None,
                metric_max=None, extra_columns=None, taskid_range=None, 
                summary=False, to_json=False, to_html=False,
                to_csv=False, to_pcp=False, to_mlflow=False,
                **task_params):
        task_params['metadata'] = True
        task_params['status'] = ['finished', 'done']
        datasets = self.taskmgr.query_tasks(**task_params)
        # filter tasks by jetitaskid range
        if taskid_range is not None:
            datasets = self.taskmgr.filter_range(datasets, taskid_range)
        col = ['computingsite', 'jobs_metadata', 'jeditaskid', 'site']
        datasets = self.taskmgr.filter_datasets(datasets, col)
        
        extra_columns = extra_columns if extra_columns is not None else []
        for extra in extra_columns:
            if not extra in kSupportedExtraColumns:
                raise ValueError('Unsupported extra column {}'.format(extra))

        skip_time = 'tims_s' not in extra_columns
        self.parse_datasets(datasets, project_name, extra_columns=extra_columns)
        self.show(extra_columns, skip_time)

        bounds = {'metric_min':metric_min, 'metric_max':metric_max}
        
        if summary:
            self.to_summary(outname+'_summary.json')
        if to_json:
            self.to_json(outname+'.json', extra_columns, skip_time, **bounds)
        if to_html:
            self.to_html(outname+'.html', extra_columns, skip_time, **bounds)
        if to_csv:
            self.to_csv(outname+'.csv', extra_columns, skip_time, **bounds)    
        if to_pcp:
            self.to_parallel_coordinate_plot(outname+'_parallel_coordinate_plot.html',
                                             extra_columns, skip_time, **bounds)   
        if to_mlflow:
            self.to_mlflow(extra_columns, skip_time, **bounds)

    def filter_invalid_datasets(self, datasets, project_name):
        filtered_datasets = []
        metadata_list = []
        for dataset in datasets:
            metadata = self.extract_metadata(dataset)
            # check if metadata is empty
            if not metadata:
                continue
            # check if it is a valid metadata generated by hpogrid
            # and if the metadata corresponds to the project of concern
            if not all( validate_job_metadata(data) and 
                       data.get('project_name', None) == project_name \
                       for data in metadata):
                continue
            filtered_datasets.append(dataset)
            metadata_list += metadata
        if len(filtered_datasets) == 0:
            print('INFO: No results found.')
            sys.exit(1)
        metrics = set([metadata['metric'] for metadata in metadata_list])
        if not len(metrics)==1 :
            raise RuntimeError('HPO results from different tasks consist of different metrics: {}'.format(metrics))
        hparams = [metadata['hyperparameters'] for metadata in metadata_list]
        if not all(hp == hparams[0] for hp in hparams):
            raise RuntimeError('HPO results from different tasks consist of different hyperparameter space')
        return filtered_datasets

    def extract_metadata(self, dataset):
        metadata_list = []
        if 'jobs_metadata' in dataset:
            jobs_metadata = dataset['jobs_metadata']
            # take care of general case where a job contains multiple sub-jobs
            # which will return a list of metadata
            if not isinstance(jobs_metadata, list):
                jobs_metadata = [jobs_metadata]
            for job_metadata in jobs_metadata:
                if not isinstance(job_metadata, dict):
                    raise ValueError('Invalid job metadata format')
                keys = list(job_metadata.keys())
                # skip empty metadata
                if not keys:
                    continue
                metadata = job_metadata[keys[0]].get('user_job_metadata', None)
                if metadata is not None:
                    metadata_list.append(metadata)
                else:
                    print('WARNING: The key "user_job_metadata" not found in job metadata. '
                          'Probably the format has changed. Skipping.')
        return metadata_list

    def parse_datasets(self, datasets, project_name, extra_columns=kDefaultExtraColumns):
        summary = []
        valid_datasets = self.filter_invalid_datasets(datasets, project_name)

        metric =  None
        mode = None
        
        for dataset in valid_datasets:
            metadata = self.extract_metadata(dataset)
            extras = {}
            if 'taskid' in extra_columns:
                extras['taskid'] = dataset['jeditaskid']
            if 'site' in extra_columns:
                if 'computingsite' in dataset:
                    extras['site'] = dataset['computingsite']
                elif 'site' in dataset:
                    extras['site'] = dataset['site']
                else:
                    extras['site'] = ''
            self.append(metadata, extras)