import argparse
import logging
import warnings

from syndatagenyolo.data_generation import (
    BlendingConfig,
    BlendingMode,
    ColorHarmonizationConfig,
    # OutputMode,
    ScalingConfig,
    SyntheticImageGenerator,
)

logger = logging.getLogger(__name__)


def main(args=None):
    parser = argparse.ArgumentParser(description="Synthetic Image Generator")
    parser.add_argument(
        "-in_dir",
        "--input_dir",
        type=str,
        required=True,
        help="Path to the input directory. It must contain a backgrounds directory and a foregrounds directory",
    )
    parser.add_argument(
        "-out_dir",
        "--output_dir",
        type=str,
        required=True,
        help="The directory where images and label files will be placed",
    )
    parser.add_argument(
        "-img_number",
        "--image_number",
        type=int,
        required=True,
        help="Number of images to create",
    )

    parser.add_argument(
        "--max_objects_per_image",
        type=int,
        default=3,
        help="Maximum number of objects per images",
    )

    parser.add_argument(
        "-i_w",
        "--image_width",
        type=int,
        default=640,
        help="Width of the output images",
    )
    parser.add_argument(
        "-i_h",
        "--image_height",
        type=int,
        default=480,
        help="Height of the output images",
    )

    parser.add_argument(
        "--augmentation_path",
        type=str,
        default="",
        required=False,
        help="Path to albumentations augmentation pipeline file",
    )

    parser.add_argument(
        "--avoid_collisions",
        default=True,
        action="store_false",
        help="Whether or not to avoid collisions (default=true)",
    )
    parser.add_argument(
        "--parallelize",
        default=False,
        action="store_true",
        help="Whether or not to use multiple cores (default=false)",
    )
    parser.add_argument(
        "--yolo_input",
        default=False,
        action="store_true",
        help="Has your background images been annotated in YOLO format?",
    )
    # output mode: either YOLO or COCO or CLASSIFICATION_SINGLE
    # parser.add_argument(
    #     "-o",
    #     "--output_mode",
    #     type=str,
    #     default="YOLO",
    #     help="Output mode: either YOLO or COCO or CLASSIFICATION_SINGLE (default=YOLO)",
    # )

    # Scaling config
    parser.add_argument(
        "--scale_foreground_by_background_size",
        default=True,
        action="store_false",
        help="Whether the foreground images should be scaled based on the background size (default=true)",
    )
    parser.add_argument(
        "-s",
        "--scaling_factors",
        type=float,
        nargs=2,
        default=(0.25, 0.85),
        help="Min and Max percentage size of the short side of the background image",
    )
    parser.add_argument(
        "--fixed_image_sizes",
        default=True,
        action="store_false",
        help="Whether or not to use fixed image sizes (default=true) => if false, the size of the background image is used and the height and width are ignored",
    )

    # Blending config
    parser.add_argument(
        "--blending_methods",
        type=str,
        nargs="+",
        default=["alpha", "gaussian", "poisson_normal", "poisson_mixed", "pyramid"],
        help="Blending methods to use (alpha, gaussian, poisson_normal, poisson_mixed, pyramid). List of strings",
    )
    parser.add_argument(
        "-p_l",
        "--pyramid_blending_levels",
        type=int,
        default=6,
        help="Number of levels for the pyramid blending method",
    )
    parser.add_argument(
        "--gaussian_options",
        type=int,
        nargs=2,
        default=(15, 30),
        help="Kernel size and sigma for Gaussian blur blending mode",
    )

    # Color harmonization config
    parser.add_argument(
        "-c",
        "--color_harmonization",
        default=False,
        action="store_true",
        help="Do you want to apply color harmonization?",
    )
    parser.add_argument(
        "-c_a",
        "--color_harmon_alpha",
        type=float,
        default=0.5,
        help="Color harmonization blending factor (0.0 to 1.0)",
    )
    parser.add_argument(
        "-c_rand",
        "--random_color_harmon_alpha",
        default=False,
        action="store_true",
        help="Randomize the color harmonization blending factor (overrides --color_harmon_alpha)",
    )

    parser.add_argument(
        "--distractor_objects",
        type=str,
        nargs="+",
        default=[],
        help="List of foreground images, which should be used as distractor objects",
    )
    parser.add_argument(
        "--overwrite_output",
        default=False,
        action="store_true",
        help="Overwrite the output directory if it exists (default=false)",
    )
    parser.add_argument(
        "--class_map_path",
        type=str,
        default="classes.yaml",
        help="Path to the YAML file containing the class names. If not provided, a default 'classes.yaml' will be used.",
    )
    parser.add_argument(
        "--debug",
        default=False,
        action="store_true",
        help="Shows the images with the yolo labels (only works with yolo_output)",
    )
    # parser.add_argument('-yolo', '--yolo_output', default=False, action='store_true',
    #                     help='Do you want to output the images in YOLO format?')
    if args is None:
        args = parser.parse_args()  # Parse args if called standalone
    else:
        args = parser.parse_args(args)  # Parse args when called from main CLI

    import sys

    blending_methods = []
    for blending_method in args.blending_methods:
        if blending_method == "alpha":
            blending_methods.append(BlendingMode.ALPHA_BLENDING)
        elif blending_method == "gaussian":
            blending_methods.append(BlendingMode.GAUSSIAN_BLUR)
        elif blending_method == "poisson_normal":
            blending_methods.append(BlendingMode.POISSON_BLENDING_NORMAL)
        elif blending_method == "poisson_mixed":
            blending_methods.append(BlendingMode.POISSON_BLENDING_MIXED)
        elif blending_method == "pyramid":
            blending_methods.append(BlendingMode.PYRAMID_BLEND)
        else:
            print(f"Blending method {blending_method} not found")
            sys.exit(1)

    # set the scaling config
    scaling_config = ScalingConfig(
        by_background_size=args.scale_foreground_by_background_size,
        factors=args.scaling_factors,
        fixed_sizes=args.fixed_image_sizes,
    )
    # set the blending config
    blending_config = BlendingConfig(
        methods=blending_methods,
        pyramid_blending_levels=args.pyramid_blending_levels,
        gaussian_kernel_size=args.gaussian_options[0],
        gaussian_sigma=args.gaussian_options[1],
    )
    # set the color harmonization config
    color_harmonization_config = ColorHarmonizationConfig(
        enabled=args.color_harmonization,
        alpha=args.color_harmon_alpha,
        random_alpha=args.random_color_harmon_alpha,
    )

    # get the output mode
    # output_mode = None
    # if args.output_mode == OutputMode.YOLO.value:
    #     output_mode = OutputMode.YOLO
    # elif args.output_mode == OutputMode.CLASSIFICATION_SINGLE.value:
    #     output_mode = OutputMode.CLASSIFICATION_SINGLE
    # elif args.output_mode == OutputMode.COCO.value:
    #     output_mode = OutputMode.COCO
    # else:
    #     print(f"Output mode {args.output_mode} not found")
    #     sys.exit(1)

    if args.fixed_image_sizes and args.yolo_input:
        warnings.warn(
            "Fixed image sizes is set to true, but yolo input is enabled. This is currently not supported!"
        )
        quit()

    # set the logging level
    logging.basicConfig(level=logging.INFO)
    if args.debug:
        logging.getLogger(__name__).setLevel(logging.DEBUG)

    if (
        args.random_color_harmon_alpha or args.color_harmon_alpha
    ) and not args.color_harmonization:
        warnings.warn(
            "Color harmonization alpha is set (either random or specific), but color harmonization is not enabled. Ignoring the alpha value"
        )

    # Divide the number of images by the number of blending methods
    img_number = int(args.image_number / len(blending_methods))
    logging.info(
        f"Generating {args.image_number} images with {img_number} images per blending method"
    )

    data_generator = SyntheticImageGenerator(
        args.input_dir,
        args.output_dir,
        img_number,
        max_objects_per_image=args.max_objects_per_image,
        image_width=args.image_width,
        image_height=args.image_height,
        augmentation_path=args.augmentation_path,
        avoid_collisions=args.avoid_collisions,
        parallelize=args.parallelize,
        yolo_input=args.yolo_input,
        # output_mode=output_mode,
        scaling_config=scaling_config,
        blending_config=blending_config,
        color_harmonization_config=color_harmonization_config,
        distractor_objects=args.distractor_objects,
        overwrite_output=args.overwrite_output,
        class_map_path=args.class_map_path,
        debug=args.debug,
    )
    data_generator.generate_images()


if __name__ == "__main__":
    # run the main function
    main()
