import argschema
import os

import pandas as pd
from pathlib import Path

import allensdk
from allensdk.brain_observatory.behavior.behavior_project_cache import \
    VisualBehaviorOphysProjectCache
from allensdk.brain_observatory.behavior.behavior_project_cache.project_metadata_writer.schemas import (  # noqa: E501
    BehaviorOphysMetadataInputSchema,
    DataReleaseToolsInputSchema
)
from allensdk.brain_observatory.data_release_utils.metadata_utils.id_generator import (  # noqa: E501
    FileIDGenerator,
)
from allensdk.brain_observatory.data_release_utils.metadata_utils.utils import (  # noqa: E501
    add_file_paths_to_metadata_table,
)

#########
# These columns should be dropped from external-facing metadata
#########
SESSION_SUPPRESS = (
    'donor_id',
    'foraging_id',
    'session_name',
    'specimen_id'
)
OPHYS_EXPERIMENTS_SUPPRESS = SESSION_SUPPRESS + (
    'behavior_session_uuid',
    'published_at',
    'isi_experiment_id'
)
OPHYS_EXPERIMENTS_SUPPRESS_FINAL = [
    'container_workflow_state',
    'experiment_workflow_state']
#########

OUTPUT_METADATA_FILENAMES = {
    'behavior_session_table': 'behavior_session_table.csv',
    'ophys_session_table': 'ophys_session_table.csv',
    'ophys_experiment_table': 'ophys_experiment_table.csv',
    'ophys_cells_table': 'ophys_cells_table.csv'
}


class BehaviorProjectMetadataWriter(argschema.ArgSchemaParser):
    """Class to write project-level metadata to csv"""

    default_schema = BehaviorOphysMetadataInputSchema
    default_output_schema = DataReleaseToolsInputSchema

    def run(self):
        """Create metadata tables and add file paths/ids.
        """
        self._initialize_metadata_writer()

        self.write_metadata()

    def _initialize_metadata_writer(self):
        """Initialize the project cache and release file information.
        """
        self._file_id_generator = FileIDGenerator()
        self._behavior_project_cache = \
            VisualBehaviorOphysProjectCache.from_lims(
                data_release_date=self.args['data_release_date'])

    def write_metadata(self):
        """Writes metadata to csv"""
        os.makedirs(self.args['output_dir'], exist_ok=True)

        self.logger.info('Writing ophys sessions table')
        self._write_ophys_sessions()
        self.logger.info('Writing ophys experiments table')
        self._write_ophys_experiments()
        self.logger.info('Writing behavior sessions table')
        self._write_behavior_sessions()
        self.logger.info('Writing ophys cells table')
        self._write_ophys_cells()

        self._write_manifest()

    def _write_behavior_sessions(
            self,
            suppress=SESSION_SUPPRESS,
            output_filename=OUTPUT_METADATA_FILENAMES[
                'behavior_session_table'],
            include_trial_metrics: bool = True
    ):
        behavior_sessions = self._behavior_project_cache. \
            get_behavior_session_table(
                suppress=suppress,
                as_df=True,
                include_trial_metrics=include_trial_metrics)

        # Add release files
        ophys_experiments = \
            self._behavior_project_cache.get_ophys_experiment_table(
                suppress=suppress, as_df=True)
        ophys_session_mask = behavior_sessions.ophys_session_id.isin(
            ophys_experiments.ophys_session_id
        )
        behavior_session_w_ophys = behavior_sessions[ophys_session_mask]
        behavior_session_w_ophys["file_id"] = \
            self._file_id_generator.dummy_value
        behavior_session_w_out_ophys = behavior_sessions[
            ~ophys_session_mask]
        behavior_session_w_out_ophys.reset_index(inplace=True)
        behavior_session_w_out_ophys = add_file_paths_to_metadata_table(
            metadata_table=behavior_session_w_out_ophys,
            id_generator=self._file_id_generator,
            file_dir=Path(self.args["behavior_nwb_dir"]),
            file_prefix=self.args["behavior_nwb_prefix"],
            index_col="behavior_session_id",
            data_dir_col="behavior_session_id",
            on_missing_file=self.args["on_missing_file"],
        )
        behavior_session_w_out_ophys.set_index("behavior_session_id",
                                               inplace=True)
        behavior_sessions = pd.concat(
            [behavior_session_w_out_ophys, behavior_session_w_ophys]
        )

        self._write_metadata_table(df=behavior_sessions,
                                   filename=output_filename)

    def _write_ophys_cells(self,
                           output_filename=OUTPUT_METADATA_FILENAMES[
                                'ophys_cells_table']):
        ophys_cells = self._behavior_project_cache. \
            get_ophys_cells_table()
        self._write_metadata_table(df=ophys_cells,
                                   filename=output_filename)

    def _write_ophys_sessions(self, suppress=SESSION_SUPPRESS,
                              output_filename=OUTPUT_METADATA_FILENAMES[
                                  'ophys_session_table'
                              ]):
        ophys_sessions = self._behavior_project_cache. \
            get_ophys_session_table(suppress=suppress, as_df=True)
        self._write_metadata_table(df=ophys_sessions,
                                   filename=output_filename)

    def _write_ophys_experiments(self, suppress=OPHYS_EXPERIMENTS_SUPPRESS,
                                 output_filename=OUTPUT_METADATA_FILENAMES[
                                     'ophys_experiment_table'
                                 ]):
        ophys_experiments = \
            self._behavior_project_cache.get_ophys_experiment_table(
                suppress=suppress, as_df=True)

        # Add release files
        ophys_experiments.reset_index(inplace=True)
        ophys_experiments = add_file_paths_to_metadata_table(
            metadata_table=ophys_experiments,
            id_generator=self._file_id_generator,
            file_dir=Path(self.args["ophys_nwb_dir"]),
            file_prefix=self.args["ophys_nwb_prefix"],
            index_col="ophys_experiment_id",
            data_dir_col="ophys_experiment_id",
            on_missing_file=self.args["on_missing_file"],
        )
        ophys_experiments.set_index('ophys_experiment_id', inplace=True)

        # users don't need to see these
        ophys_experiments.drop(
                labels=OPHYS_EXPERIMENTS_SUPPRESS_FINAL,
                inplace=True,
                axis=1)

        self._write_metadata_table(df=ophys_experiments,
                                   filename=output_filename)

        return ophys_experiments

    def _write_metadata_table(self, df: pd.DataFrame, filename: str):
        """
        Writes file to csv

        Parameters
        ----------
        df
            The dataframe to write
        filename
            Filename to save as
        """
        filepath = os.path.join(self.args["output_dir"], filename)

        self.logger.info(f'Writing {filepath}')

        df = df.reset_index()
        df.to_csv(filepath, index=False)

    def _write_manifest(self):
        def get_abs_path(filename):
            return os.path.abspath(os.path.join(self.args["output_dir"],
                                                filename))

        metadata_filenames = OUTPUT_METADATA_FILENAMES.values()
        metadata_files = [get_abs_path(f) for f in metadata_filenames]
        pipeline_metadata = []
        sdk_metadata = {
            "name": "AllenSDK",
            "version": str(allensdk.__version__),
            "comment": "",
        }
        pipeline_metadata.append(sdk_metadata)
        output_data = {
            "metadata_files": metadata_files,
            "data_pipeline_metadata": pipeline_metadata,
            "project_name": "visual-behavior-ophys",
            "log_level": "INFO",
        }

        self.output(output_data, indent=2)
