import time
from pathlib import Path
import typer
from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn

app = typer.Typer()
console = Console()
table = Table(title="Analysis Results 🥳")


@app.command()
def sdh_analysis(
    image_path: Path = typer.Argument(
        ...,
        help="The image file path to analyse.",
        exists=True,
        file_okay=True,
        dir_okay=False,
        writable=False,
        readable=True,
        resolve_path=True,
    ),
    mask_path: Path = typer.Option(
        None,
        help="The path to a binary mask to hide slide region during analysis. It needs to be of the same resolution as input image and only pixel marked as 1 will be analyzed.",
        exists=True,
        file_okay=True,
        dir_okay=False,
        writable=False,
        readable=True,
        resolve_path=True,
    ),
    model_path: Path = typer.Option(
        None,
        help="The SDH model path to use for analysis. Will download latest one if no path provided.",
        exists=True,
        file_okay=True,
        dir_okay=False,
        writable=False,
        readable=True,
        resolve_path=True,
    ),
    cellpose_path: Path = typer.Option(
        None,
        help="The pre-computed CellPose mask to use for analysis. Will run Cellpose if no path provided. Required as an image file.",
        exists=True,
        file_okay=True,
        dir_okay=False,
        writable=False,
        readable=True,
        resolve_path=True,
    ),
    output_path: Path = typer.Option(
        None,
        help="The path to the folder to save the results. Will save in the current folder if not specified.",
    ),
    cellpose_diameter: int = typer.Option(
        None,
        help="Approximative single cell diameter in pixel for CellPose detection. If not specified, Cellpose will try to deduce it.",
    ),
    export_map: bool = typer.Option(
        True,
        help="Export the original image with cells painted by classification label.",
    ),
    export_stats: bool = typer.Option(True, help="Export per fiber stat table."),
):
    """Run the SDH analysis and quantification on the image."""
    start_time = time.time()
    console.print(
        f"👋 [bold dark_orange]Welcome to the mitochondrial distribution analysis (SDH images)",
    )
    console.print(f"📄 INPUT: raw image: {image_path}", style="blue")

    if cellpose_path is None:
        console.print(
            "💡 INFO: No CellPose mask provided, will run CellPose during the analysis.",
            style="blue",
        )
    else:
        console.print(f"📄 INPUT: CellPose mask: {cellpose_path}", style="blue")

    if mask_path is not None:
        console.print(f"📄 INPUT: binary mask: {mask_path}", style="blue")
    import os

    if model_path is None:
        console.print(
            "💡 INFO: No SDH model provided, will download or use latest one.",
            style="blue",
        )
        model_path_abs = Path(os.path.abspath(__file__)).parents[0] / "model.h5"
        if not os.path.exists(model_path_abs):
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                TimeElapsedColumn(),
                transient=False,
            ) as progress:
                progress.add_task(description="Downloading SDH Model...", total=None)
                urllib.request.urlretrieve(
                    "https://lbgi.fr/~meyer/SDH_models/model.h5",
                    model_path_abs,
                )
        model_path = model_path_abs
    console.print(f"📄 INPUT: SDH Model: {model_path}", style="blue")

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TimeElapsedColumn(),
        transient=False,
    ) as progress:
        progress.add_task(description="Importing the libraries...", total=None)
        import urllib.request
        from ..src.common_func import (
            is_gpu_availiable,
            load_cellpose,
            load_sdh_model,
            run_cellpose,
            label2rgb,
            blend_image_with_label,
            HiddenPrints,
        )
        from ..src.SDH_analysis import run_sdh_analysis
        import numpy as np
        from PIL import Image

        try:
            from imageio.v2 import imread
        except:
            from imageio import imread

    if output_path is None:
        output_path = image_path.parents[0]
    else:
        Path(output_path).mkdir(parents=True, exist_ok=True)

    if is_gpu_availiable():
        console.print(f"💡 INFO: GPU is available.", style="blue")
    else:
        console.print(f"❌ INFO: GPU is not available. Using CPU only.", style="red")

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TimeElapsedColumn(),
        transient=False,
    ) as progress:
        progress.add_task(description="Reading all inputs...", total=None)
        image_ndarray_sdh = imread(image_path)

    if mask_path is not None:
        mask_ndarray = imread(mask_path)
        if np.unique(mask_ndarray).shape[0] != 2:
            console.print(
                "The mask image should be a binary image with only 2 values (0 and 1).",
                style="red",
            )
            raise ValueError
        if len(image_ndarray_sdh.shape) > 2:
            mask_ndarray = np.repeat(
                mask_ndarray.reshape(mask_ndarray.shape[0], mask_ndarray.shape[1], 1),
                image_ndarray_sdh.shape[2],
                axis=2,
            )
        image_ndarray_sdh = image_ndarray_sdh * mask_ndarray
        if cellpose_path is not None:
            mask_cellpose = imread(cellpose_path)

    if cellpose_path is None:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            TimeElapsedColumn(),
            transient=False,
        ) as progress:
            progress.add_task(description="Running CellPose...", total=None)
            model_cellpose = load_cellpose()
            mask_cellpose = run_cellpose(
                image_ndarray_sdh, model_cellpose, cellpose_diameter
            )
            mask_cellpose = mask_cellpose.astype(np.uint16)
            cellpose_mask_filename = image_path.stem + "_cellpose_mask.tiff"
            Image.fromarray(mask_cellpose).save(output_path / cellpose_mask_filename)
        console.print(
            f"💾 OUTPUT: CellPose mask saved as {output_path/cellpose_mask_filename}",
            style="green",
        )

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TimeElapsedColumn(),
        transient=False,
    ) as progress:
        progress.add_task(
            description="Loading SDH Model...",
            total=None,
        )
        with HiddenPrints():
            model_SDH = load_sdh_model(model_path)

    if mask_path is not None:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            TimeElapsedColumn(),
            transient=False,
        ) as progress:
            progress.add_task(
                description="Masking Cellpose and Stardist mask with binary mask...",
                total=None,
            )
        mask_ndarray = imread(mask_path)
        mask_cellpose = mask_cellpose * mask_ndarray

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TimeElapsedColumn(),
        transient=False,
    ) as progress:
        progress.add_task(
            description="Predicting all muscle fibers class...", total=None
        )
        with HiddenPrints():
            result_df, full_label_map, df_cellpose_details = run_sdh_analysis(
                image_ndarray_sdh, model_SDH, mask_cellpose
            )
    if export_map:
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            TimeElapsedColumn(),
            transient=False,
        ) as progress:
            progress.add_task(
                description="Blending label and original image together...", total=None
            )
            labelRGB_map = label2rgb(image_ndarray_sdh, full_label_map)
            overlay_img = blend_image_with_label(image_ndarray_sdh, labelRGB_map)
            overlay_filename = image_path.stem + "_label_blend.tiff"
            overlay_img.save(output_path / overlay_filename)

    table.add_column("Feature", justify="left", style="cyan")
    table.add_column("Raw Count", justify="center", style="magenta")
    table.add_column("Proportion (%)", justify="right", style="green")
    for index, row in result_df.iterrows():
        table.add_row(
            str(row[0]),
            str(row[1]),
            str(row[2]),
        )
    console.print(table)
    csv_name = image_path.stem + "_results.csv"
    result_df.to_csv(
        output_path / csv_name,
        index=False,
    )
    console.print(
        f"💾 OUTPUT: Summary Table saved as {output_path/csv_name}", style="green"
    )
    if export_stats:
        cell_details_name = image_path.stem + "_cell_details.csv"
        df_cellpose_details.drop("image", axis=1).to_csv(
            output_path / cell_details_name,
            index=False,
        )
        console.print(
            f"💾 OUTPUT: Cell Table saved as {output_path/cell_details_name}",
            style="green",
        )
    label_map_name = image_path.stem + "_label_map.tiff"
    Image.fromarray(full_label_map).save(output_path / label_map_name)
    console.print(
        f"💾 OUTPUT: Label map saved as {output_path/label_map_name}", style="green"
    )
    console.print("--- %s seconds ---" % int((time.time() - start_time)))
