import argschema.fields
from allensdk.brain_observatory.argschema_utilities import RaisingSchema
from allensdk.brain_observatory.ecephys.write_nwb.schemas import (
    BaseBehaviorSessionDataSchema,
    BaseNeuropixelsSchema,
)
from argschema import ArgSchema
from argschema.fields import LogLevel


class _VBNSessionDataSchema(
    BaseBehaviorSessionDataSchema, BaseNeuropixelsSchema
):
    mapping_stimulus_file = argschema.fields.InputFile(
        required=True, description="path to mapping_stimulus_file"
    )
    replay_stimulus_file = argschema.fields.InputFile(
        required=True, description="path to replay_stimulus_file"
    )
    stim_table_file = argschema.fields.InputFile(
        required=True, description="path to stimulus presentations csv file"
    )
    raw_eye_tracking_video_meta_data = argschema.fields.InputFile(
        required=True, description="path to eye tracking metadata"
    )
    eye_dlc_file = argschema.fields.InputFile(
        required=True, description="path to deeplabcut eye tracking h5 file"
    )
    side_dlc_file = argschema.fields.InputFile(
        required=True, description="path to deeplabcut side tracking h5 file"
    )
    face_dlc_file = argschema.fields.InputFile(
        required=True, description="path to deeplabcut face tracking h5 file"
    )
    eye_tracking_filepath = argschema.fields.InputFile(
        required=True,
        description="h5 filepath containing eye tracking ellipses",
    )
    sync_file = argschema.fields.InputFile(
        required=True, description="path to sync file"
    )
    ecephys_session_id = argschema.fields.Int(
        required=True, description="ecephys session id"
    )


class VBNInputSchema(ArgSchema):
    """Input schema for visual behavior neuropixels"""

    log_level = LogLevel(
        default="INFO", description="Logging level of the module"
    )
    session_data = argschema.fields.Nested(
        _VBNSessionDataSchema,
        required=True,
        description="Data pertaining to a behavior session",
    )
    skip_probes = argschema.fields.List(
        argschema.fields.Str,
        cli_as_single_argument=True,
        default=None,
        allow_none=True,
    )
    output_path = argschema.fields.OutputFile(
        required=True, description="Path of output.json to be written"
    )


class OutputSchema(RaisingSchema):
    input_parameters = argschema.fields.Nested(VBNInputSchema)
    output_path = argschema.fields.OutputFile(
        required=True, description="write outputs to here"
    )
