"""Run the main code for Deep-Package-Detection"""

from pathlib import Path
import logging

import click

from deep_package_detection import __version__
from deep_package_detection.loggings import config_logger
from deep_package_detection.detector import PackageDetectorTrainer
from deep_package_detection.detector import PackageDetectorInference


logger = logging.getLogger(__name__)


@click.group()
@click.version_option(version=__version__)
@click.option("-l", "--log_path", type=str, help="Path to save log file")
@click.option(
    "-v",
    "--verbose",
    count=True,
    help="Shorthand for info/debug/warning/error loglevel (-v/-vv/-vvv/-vvvv)",
)
def deep_package_detection_cli(log_path: str, verbose: int) -> None:
    """This is a deep model for detecting different types of packages."""
    if verbose == 1:
        log_level = 10
    elif verbose == 2:
        log_level = 20
    elif verbose == 3:
        log_level = 30
    else:
        log_level = 40
    config_logger(log_level, log_path)


@deep_package_detection_cli.command()
@click.option("--img_resize", type=int, default=640, help="Resize images to this size.")
@click.option(
    "--conf_path",
    type=str,
    default="src/deep_package_detection/data/data.yaml",
    help="Path to the config file",
)
@click.option(
    "--epochs", type=int, default=100, help="Number of epochs used in training."
)
@click.option("--batch_size", type=int, default=16, help="Batch size used in training.")
@click.option(
    "--device", type=str, default="cuda", help="Use cuda or cpu for training."
)
def train(
    img_resize: int, conf_path: str, epochs: int, batch_size: int, device: str
) -> None:
    """This the CLI for training purposes"""
    detector = PackageDetectorTrainer(
        conf=conf_path,
        img_size=img_resize,
        epochs=epochs,
        device=device,
        batch_size=batch_size,
    )
    detector.train()
    _ = detector.validation()


@deep_package_detection_cli.command()
@click.option(
    "--model_path",
    type=click.Path(),
    default=Path("./src/deep_package_detection/model/package_detection.pt"),
    help="Path to the pre-trained model.",
)
@click.option(
    "--data_path",
    type=click.Path(),
    default=Path("./tests/test_data"),
    help="Path to the test data.",
)
@click.option(
    "--result_path", type=str, default="./results", help="Path to the results."
)
def infer(model_path: Path, data_path: str, result_path: str) -> None:
    """This the CLI for testing purposes"""
    logger.info("Testing the model for package detection...")
    inferer = PackageDetectorInference(
        model_path=Path(model_path), result_path=result_path
    )
    segmentations = inferer.inference(data_path=data_path)
    counts = inferer.count_packages(segmentations)
    if counts:
        for key, val in counts.items():
            logger.info(
                "%s packages are detected in %s as: %s",
                sum(item["count"] for item in val),
                key,
                val,
            )
    inferer.plot_and_save_results(segmentations)