"""Configure Slideflow projects for reproducing published results."""

import os
import click
import slideflow as sf
from biscuit import experiment
from os.path import exists, join, abspath

# -----------------------------------------------------------------------------

@click.command()
@click.option('--train_slides', type=str, help='Directory to training slides, for cross-validation', required=True)
@click.option('--train_anns', type=str, help='Directory to annotation file for training data (CSV)', default='annotations/tcga.csv', show_default=True)
@click.option('--train_roi', type=str, help='Directory to CSV ROI files, for cross-validation')
@click.option('--outcome', type=str, help='Outcome (annotation header) that assigns class labels.', default='cohort', show_default=True)
@click.option('--outcome1', type=str, help='First class label.', default='LUAD', show_default=True)
@click.option('--outcome2', type=str, help='Second class label.', default='LUSC', show_default=True)
@click.option('--val_slides', type=str, help='Directory to external evaluation slides, for evaluation')
@click.option('--val_anns', type=str, help='Directory to annotation file for training data (CSV)', default='annotations/cptac.csv', show_default=True)
def configure_projects(
    train_slides,
    train_anns,
    train_roi,
    outcome,
    outcome1,
    outcome2,
    val_slides=None,
    val_anns=None,
):
    """Configure Slideflow projects for reproducing published results.

    This script uses the provided slides to build Slideflow projects in the
    'projects/' folder of the current working directory. Clinical annotations
    (class labels) are read from the 'annotations/' folder unless otherwise
    specified.

    Training slides from The Cancer Genome Atlas (TCGA) are available at
    https://portal.gdc.cancer.gov/ (projects TCGA-LUAD and TCGA-LUSC).

    Validation slides from the Clinical Proteomics Tumor Analysis Consortium
    (CPTAC) are available at https://proteomics.cancer.gov/data-portal
    (projects CPTAC-LUAD and CPTAC-LSCC).
    """

    # Absolute paths
    train_slides = abspath(train_slides)
    train_anns = abspath(train_anns)
    out = abspath('projects')
    if val_slides:
        val_slides = abspath(train_slides)
    if val_anns:
        val_anns = abspath(val_anns)
    if train_roi:
        train_roi = abspath(train_roi)
    gan_path = abspath('gan')
    if not exists(gan_path):
        os.makedirs(gan_path)

    # --- Set up projects -----------------------------------------------------

    # Set up training project
    if (not exists(join(out, 'training'))
       or not exists(join(out, 'training', 'settings.json'))):
        print("Setting up training project...")
        tP = sf.Project(
            join(out, 'training'),
            sources=['Training'],
            annotations=train_anns
        )
        tP.add_source(
            name='Training',
            slides=train_slides,
            roi=(train_roi if train_roi else train_slides),
            tiles=join(out, 'training', 'tiles'),
            tfrecords=join(out, 'training', 'tfrecords')
        )
        tP.add_source(
            name='LUNG_GAN',
            slides=gan_path,
            roi=gan_path,
            tiles=gan_path,
            tfrecords=gan_path
        )
        print(f"Training project setup at {join(out, 'training')}.")
    else:
        tP = sf.Project(join(out, 'training'))
        print("Loading training project which already exists.")

    # Set up external evaluation project
    if val_slides:
        if not val_anns:
            msg = "If providing evaluation slides, evaluation annotations "
            msg += "must also be provided (--val_anns)"
            raise ValueError(msg)
        if (not exists(join(out, 'evaluation'))
           or not exists(join(out, 'evaluation', 'settings.json'))):
            print("Setting up evaluation project.")
            eP = sf.Project(
                join(out, 'evaluation'),
                sources=['Evaluation'],
                annotations=val_anns
            )
            eP.add_source(
                name='Evaluation',
                slides=val_slides,
                roi=val_slides,
                tiles=join(out, 'evaluation', 'tiles'),
                tfrecords=join(out, 'evaluation', 'tfrecords')
            )
            print(f"Evaluation project setup at {join(out, 'evaluation')}.")
        else:
            eP = sf.Project(join(out, 'evaluation'))
            print("Loading evaluation project which already exists.")

    # --- Perform tile extraction ---------------------------------------------

    print("Extracting tiles from WSIs at 299px, 302um")
    for P in (eP, tP):
        P.extract_tiles(
            tile_px=299,
            tile_um=302,
            qc='both',
            img_format='png'
        )
    print("Extracting tiles from WSIs at 512px, 400um (for GAN training)")
    for P in (eP, tP):
        P.extract_tiles(
            tile_px=512,
            tile_um=400,
            qc='both',
            img_format='png'
        )
    print("Finished tile extraction, project configuration complete.")

    # --- Save GAN training configuration -------------------------------------

    if not exists('gan_config.json'):
        gan_config = {
            "project_path": join(out, 'training'),
            "tile_px": 512,
            "tile_um": 400,
            "model_type": "categorical",
            "outcomes": [outcome],
            "filters": {outcome: [outcome1, outcome2]}
        }
        sf.util.write_json(gan_config, 'gan_config.json')
        print("Wrote GAN configuration to gan_config.json")
    else:
        print("GAN configuration already exists at gan_config.json")


# ----------------------------------------------------------------------------

if __name__ == "__main__":
    configure_projects()  # pylint: disable=no-value-for-parameter

# ----------------------------------------------------------------------------