import argparse
import json
from typing import Any
import numpy as np
from pathlib import Path
import re
import subprocess
import time
import sys


from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.core.termination import NoTermination
from pymoo.core.evaluator import Evaluator
from pymoo.core.population import Population
from pymoo.problems.static import StaticProblem
from pymoo.operators.crossover.pntx import PointCrossover
from pymoo.operators.mutation.bitflip import BitflipMutation

from editing_DiTs.benchmark.score_images import DEFAULT_BENCHMARK_PROMPTS
from editing_DiTs.genetic.flux_population_io_manager import (
    FluxPopulationIOManager,
)
from editing_DiTs.genetic.flux_problem import FluxCachingScheduleProblem
from editing_DiTs.genetic.pixart_population_io_manager import (
    PixArtPopulationIOManager,
)
from editing_DiTs.genetic.pixart_problem import PixArtCachingScheduleProblem
from editing_DiTs.genetic.sampling import BinaryRandomSampling
from editing_DiTs.genetic.population_io_manager import (
    PopulationIOManager,
    DEFAULT_POPULATIONS_DIR,
    DEFAULT_BENCHMARKS_DIR,
)
from editing_DiTs.image_generators.image_generator import ImageGenerator
from editing_DiTs.image_generators.load_image_generator import (
    ImageGeneratorRegistry,
    get_image_generator_type,
)
from editing_DiTs.image_generators.pixart_image_generator import (
    PixArtImageGenerator,
)
from editing_DiTs.image_generators.flux_image_generator import (
    FluxImageGenerator,
)


def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "Run n cycles of NSGA-II training with offline evaluation."
            "Either provide a path to load an existing population manager or"
            "specify all other arguments to start a new run."
        )
    )
    parser.add_argument(
        "--image-generator",
        type=str,
        required=True,
        help="The name of the image generator to use.",
        choices=list(ImageGeneratorRegistry.registry.keys()),
    )
    parser.add_argument(
        "--load-from",
        type=Path,
        help="Load an existing population manager from a directory. If not provided, all other arguments are required.",
    )
    parser.add_argument(
        "--name", type=str, help="Name of the population manager"
    )
    parser.add_argument(
        "--num-cycles",
        default="1",
        type=str,
        help="Number of cycles to run (default: 1). If `inf` is provided, run indefinitely.",
    )
    parser.add_argument(
        "--print-not-submit",
        action="store_true",
        help="Print commands to run offline evaluation scripts without submitting them. Num cycles must be 1.",
    )
    parser.add_argument(
        "--all-populations-dir",
        type=Path,
        default=DEFAULT_POPULATIONS_DIR,
        help="Base directory for population schedules",
    )
    parser.add_argument(
        "--all-benchmarks-dir",
        type=Path,
        default=DEFAULT_BENCHMARKS_DIR,
        help="Base directory for benchmark scores",
    )
    parser.add_argument(
        "--embedding-dir",
        type=Path,
        required=True,
        help="The direcotry containing embeddings used to generate images",
    )
    parser.add_argument(
        "--generation-num",
        type=int,
        default=None,
        help="Generation number to start from (default: use highest existing or 0)",
    )
    parser.add_argument(
        "--num-inference-steps",
        type=int,
        default=20,
        help="Number of inference steps",
    )
    parser.add_argument(
        "--benchmark-prompts",
        type=Path,
        required=True,
        help="JSON file containing benchmark prompts.",
    )
    parser.add_argument(
        "--file-mode",
        choices=["json", "text", "txt", "tsv"],
        help="For scoring images. JSON is the Image Reward prompts mode. TSV must have a column named 'Prompt', and have a header line. Text format is a simple text file with a prompt on each line.",
    )
    parser.add_argument(
        "--image-naming-mode",
        choices=["image_reward", "parti"],
        help="The pattern to use to match image filenames. Use image_reward for image_reward embeddings, parti for parti and drawbench embeddings.",
    )
    parser.add_argument(
        "--min-diff-from-default",
        type=int,
        default=1,
        help="Minimum number of differences required from default schedule",
    )
    parser.add_argument(
        "--population-size",
        type=int,
        default=72,
        help="Population size for NSGA-II",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=100,
        help="Batch size for image generation runs",
    )
    parser.add_argument(
        "--num-images-per-prompt",
        type=int,
        default=10,
        help="Number of images to generate per prompt.",
    )
    parser.add_argument(
        "--exactly-n-images",
        type=int,
        default=1000,
        help="Number of images to expect before scoring them",
    )
    parser.add_argument(
        "--regen-if-not-n-images",
        type=int,
        default=1000,
        help="Force re-generating images for schedules whose directories do not contain exactly n images.",
    )
    parser.add_argument(
        "--maximize-MACs",
        action="store_true",
        help="Maximize MACs instead of minimizing them.",
    )
    parser.add_argument(
        "--qos",
        type=str,
        default="REDACTED",
        help="The QOS to use for offline evaulations.",
    )
    parser.add_argument(
        "--gpu-generate",
        type=str,
        default="rtxa5000",
        help="GPU to use for image generation runs",
    )
    parser.add_argument(
        "--gpu-compute",
        type=str,
        default="rtxa5000",
        help="GPU to use for MAC compute runs",
    )
    parser.add_argument(
        "--mem-generate",
        type=str,
        default="32",
        help="Memory to use for image generation runs",
    )
    parser.add_argument(
        "--mem-compute",
        type=str,
        default="32",
        help="Memory to use for MAC compute runs",
    )
    parser.add_argument(
        "--time-generate",
        type=str,
        default="00:07:00",
        help="Time to allocate for image generation runs",
    )
    parser.add_argument(
        "--time-compute",
        type=str,
        default="00:07:00",
        help="Time to allocate for MAC compute runs",
    )

    args = parser.parse_args()

    if args.load_from is None and args.name is None:
        raise ValueError("Must provide a name or path to JSON file.")

    if args.num_cycles.isdigit():
        args.num_cycles = int(args.num_cycles)
    elif args.num_cycles != "inf":
        raise ValueError("num_cycles must be an integer or 'inf'.")

    if args.num_cycles != 1 and args.print_not_submit:
        raise ValueError(
            "print-not-submit can only be used with num_cycles=1."
        )

    return args


def init_gen_0(
    manager: PopulationIOManager, image_generator_type: type[ImageGenerator]
) -> NSGA2:
    print(
        "Attempting to load seed population. Should be under gen_000/candidates/cand_XXX.json"
    )
    sample = False
    X = manager.load_population_vectors(0)

    if X.size == 0:
        if (
            input(
                "No inital population found. Randomly seed the population? (y/N) "
            )
            == "y"
        ):
            sample = True
        else:
            raise ValueError("No initial population found")

    print("Initializing new algorithm.")
    if issubclass(image_generator_type, PixArtImageGenerator):
        assert isinstance(manager, PixArtPopulationIOManager)
        problem = PixArtCachingScheduleProblem(
            manager.num_inference_steps,
            manager.num_blocks,
            manager.num_component_types,
            manager.min_diff_from_default,
        )
    elif issubclass(image_generator_type, FluxImageGenerator):
        assert isinstance(manager, FluxPopulationIOManager)
        problem = FluxCachingScheduleProblem(
            manager.num_inference_steps,
            manager.num_blocks,
            manager.num_single_blocks,
            manager.num_component_types_full,
            manager.num_component_types_single,
            manager.min_diff_from_default,
        )
    else:
        raise ValueError("Unsupported image generator type.")

    if sample:
        sampling_or_pop = BinaryRandomSampling()
    else:
        sampling_or_pop = Population.new("X", X)

    algorithm = NSGA2(
        pop_size=manager.population_size,
        sampling=sampling_or_pop,
        crossover=PointCrossover(prob=0.9, n_points=4),
        mutation=BitflipMutation(prob=0.05),
    )
    print("Setting up algorithm.")
    termination = NoTermination()
    algorithm.setup(problem, termination=termination, seed=0, verbose=True)
    init_pop = algorithm.ask()

    # need to sync, since algo now has gen 1, while manager has gen 0 for seed gen
    manager.generation_num = algorithm.n_gen
    X = init_pop.get("X")

    manager.save_population(X, 1)
    print(f"Candidates for generation {manager.generation_num} saved. ")

    return algorithm


def train_one_cycle(manager: PopulationIOManager, algorithm: NSGA2) -> bool:
    if not manager.check_offline_eval():
        print(
            f"Offline evaluation not found for generation {manager.generation_num}. "
        )
        return False
    else:
        # Offline evaluation is complete; load evaluation results.
        print("Offline evaluation detected. Loading evaluation results...")
        X, F, G = manager.ask()
        sampling_or_pop = Population.new("X", X, F=F, G=G)
        algorithm.tell(sampling_or_pop)

        # Sync generation number between algorithm and IO manager.
        manager.generation_num = algorithm.n_gen

        sampling_or_pop = algorithm.ask()
        X = sampling_or_pop.get("X")
        manager.save_population(X)
        print(
            f"Candidates for generation {manager.generation_num} saved. "
            "Please run your offline evaluation scripts and then re-run this script to resume."
            "Commands to run the offline evaluation are printed below."
        )
        return True


def checkpoint_manager_and_algorithm(
    manager: PopulationIOManager, algorithm: NSGA2
):
    manager.save_algorithm(algorithm)
    manager.to_json()
    print(
        f"Checkpointed manager config and algorithim for generation {manager.generation_num}"
    )


def train_nsga2(
    manager: PopulationIOManager,
    algorithm: NSGA2,
    image_generator: str,
    num_cycles: int | None,
    print_not_submit: bool = True,
    batch_size: int | None = None,
    offline_eval_kwargs: dict[str, Any] | None = None,
):
    if offline_eval_kwargs is None:
        offline_eval_kwargs = {}

    failed = False
    while num_cycles is None or num_cycles > 0:
        step_taken = train_one_cycle(manager, algorithm)
        checkpoint_manager_and_algorithm(manager, algorithm)

        if not print_not_submit and not step_taken:
            failed = False

        gen_images_cmd, score_images_cmd, compute_metrics_cmd = (
            get_offline_eval_commands_from_manager(
                manager,
                image_generator,
                print_commands=print_not_submit,
                batch_size=batch_size,
                **offline_eval_kwargs,
            )
        )

        if not print_not_submit:
            submit_and_monitor_job([gen_images_cmd], retry_count=5)
            print("-" * 80 + "\nFinished generating images.\n" + "-" * 80)

            # score_images and compute_metrics can happen at same time
            submit_and_monitor_job([score_images_cmd, compute_metrics_cmd])
            print(
                "-" * 80
                + "\nFinished scoring images and computing metrics.\n"
                + "-" * 80
            )

        if num_cycles is not None:
            num_cycles -= 1

        if failed:
            raise ValueError("Offline evaluation not found.")


def main(args):
    if args.load_from is None and args.name is None:
        raise ValueError("Must provide a name or path to JSON file.")
    image_generator_type = get_image_generator_type(args.image_generator)

    if issubclass(image_generator_type, PixArtImageGenerator):
        manager_type = PixArtPopulationIOManager
    elif issubclass(image_generator_type, FluxImageGenerator):
        manager_type = FluxPopulationIOManager
    else:
        raise ValueError(
            f"Unsupported Image Generator type {image_generator_type}."
        )

    if args.load_from is not None:
        if not args.load_from.exists():
            raise ValueError(f"JSON file {args.load_from} does not exist.")

        manager = manager_type.from_json(args.load_from)
    else:
        # Initialize the IO manager with the provided arguments.
        manager = manager_type(
            name=args.name,
            all_populations_dir=args.all_populations_dir,
            all_benchmarks_dir=args.all_benchmarks_dir,
            generation_num=args.generation_num,
            num_inference_steps=args.num_inference_steps,
            min_diff_from_default=args.min_diff_from_default,
            population_size=args.population_size,
            maximize_macs=args.maximize_MACs,
        )

    try:
        print("Attempting to load from checkpoint...")
        algorithm = manager.load_algorithm()
    except FileNotFoundError as e:
        print(f"Checkpoint not found: {e}")
        algorithm = init_gen_0(manager, image_generator_type)

    offline_eval_kwargs = {
        "embedding_dir": args.embedding_dir,
        "gpu_generate": args.gpu_generate,
        "gpu_compute": args.gpu_compute,
        "mem_generate": args.mem_generate,
        "mem_compute": args.mem_compute,
        "time_generate": args.time_generate,
        "time_compute": args.time_compute,
        "num_images_per_prompt": args.num_images_per_prompt,
        "exactly_n_images": args.exactly_n_images,
        "regen_if_not_n_images": args.regen_if_not_n_images,
        "qos": args.qos,
        "benchmark_prompts": args.benchmark_prompts,
        "file_mode": args.file_mode,
        "image_naming_mode": args.image_naming_mode,
    }

    num_cycles = None if args.num_cycles == "inf" else args.num_cycles
    train_nsga2(
        manager,
        algorithm,
        args.image_generator,
        num_cycles,
        args.print_not_submit,
        args.batch_size,
        offline_eval_kwargs=offline_eval_kwargs,
    )


def get_offline_eval_commands_from_manager(
    manager: PopulationIOManager,
    image_generator: str,
    multi_job_benchmark_scripts_base: Path = Path(
        "./results/multi-job-benchmark-scripts/genetic"
    ),
    multi_job_score_scripts_base: Path = Path(
        "results/multi-job-score-scripts/genetic"
    ),
    multi_job_compute_macs_scripts_base: Path = Path(
        "results/multi-job-compute-macs-scripts/genetic"
    ),
    gpu_generate: str = "rtxa5000",
    gpu_score: str = "rtxa5000",
    gpu_compute: str = "rtxa4000",
    time_generate: str = "00:07:00",
    time_score: str = "00:07:00",
    time_compute: str = "00:04:00",
    regen_if_not_n_images: int = 1000,
    exactly_n_images: int = 1000,
    qos: str = "REDACTED",
    mem_generate: int = 32,
    mem_score: int = 20,
    mem_compute: int = 32,
    print_commands: bool = True,
    batch_size: int | None = None,
    embedding_dir: Path | None = None,
    num_images_per_prompt: int | None = None,
    benchmark_prompts: Path | None = None,
    file_mode: str | None = None,
    image_naming_mode: str | None = None,
) -> tuple[str, str, str]:
    """
    Given a PopulationIOManager instance, print out the offline evaluation commands
    for:
      - Generating images for a generation.
      - Scoring the generated images.
      - Computing metrics from the candidates.

    The function uses the manager's name, generation number, and directory structure.
    """
    # Resolve the base script directories, to standardize user input
    benchmark_scripts_base = multi_job_benchmark_scripts_base.resolve()
    score_scripts_base = multi_job_score_scripts_base.resolve()
    compute_macs_scripts_base = multi_job_compute_macs_scripts_base.resolve()

    # Build the generation folder string, e.g., "gen_005".
    generation_str = f"gen_{manager.generation_num:03d}/candidates"

    # Construct directories for each command.
    scripts_dir_generate = (
        benchmark_scripts_base / manager.name / generation_str
    ).relative_to(Path.cwd())
    schedule_dir = manager.get_pop_candidates_dir()
    image_dir = manager.get_benchmark_gen_dir()

    scripts_dir_score = (
        score_scripts_base / manager.name / generation_str
    ).relative_to(Path.cwd())
    scripts_dir_compute = (
        compute_macs_scripts_base / manager.name / generation_str
    ).relative_to(Path.cwd())

    batch_size_str = (
        f"--batch-size {batch_size} " if batch_size is not None else ""
    )
    embedding_dir_str = (
        f"--embedding-dir {embedding_dir} "
        if embedding_dir is not None
        else ""
    )
    num_images_per_prompt_str = (
        f"--num-images-per-prompt {num_images_per_prompt} "
        if num_images_per_prompt is not None
        else ""
    )

    # Build command for generating images.
    gen_images_cmd = (
        f"python editing_DiTs/benchmark/generate_images_multi_job.py "
        f"--image-generator {image_generator} "
        f"--scripts-dir {scripts_dir_generate} "
        f"--schedule-dir {schedule_dir} "
        f"--image-dir {image_dir} "
        f"--qos {qos} "
        f"--time-per-job {time_generate} "
        f"--regen-if-not-n-images {regen_if_not_n_images} "
        f'--gpu "{gpu_generate}" '
        f"--mem {mem_generate} "
        f"{batch_size_str} "
        f"{embedding_dir_str} "
        f"{num_images_per_prompt_str} "
    )

    # Build command for scoring images.
    file_mode_str = f"--file-mode {file_mode}" if file_mode is not None else ""
    image_naming_mode = f"--image-naming-mode {image_naming_mode}" if image_naming_mode is not None else ""
    benchmark_prompts_str = f"--benchmark-prompts {benchmark_prompts}" if benchmark_prompts is not None else ""
    score_images_cmd = (
        f"python editing_DiTs/benchmark/score_images_multi_job.py "
        f'--time-per-job "{time_score}" '
        f"--gpu {gpu_score} "
        f"--mem {mem_score} "
        f"--scripts-dir {scripts_dir_score} "
        f"--image-dir {image_dir} "
        f"--exactly-n-images {exactly_n_images} "
        f"--qos {qos} "
        f"--delete-after "
        f" {benchmark_prompts_str} {file_mode_str} {image_naming_mode} "
    )

    # Build command for computing metrics.
    compute_metrics_cmd = (
        f"python editing_DiTs/benchmark/compute_macs_multi_job.py "
        f"--image-generator {image_generator} "
        f'--time-per-job "{time_compute}" '
        f'--gpu "{gpu_compute}" '
        f"--mem {mem_compute} "
        f"--qos {qos} "
        f"--scripts-dir {scripts_dir_compute} "
        f"--input-dir {schedule_dir}"
    )

    if print_commands:
        print("Command to run generation:")
        print(gen_images_cmd)
        print("\nCommand to score images:")
        print(score_images_cmd)
        print("\nCommand to compute metrics:")
        print(compute_metrics_cmd)

    return gen_images_cmd, score_images_cmd, compute_metrics_cmd


def submit_job(command: str) -> int:
    """
    Submit a job using the provided command string.
    The command's output must include a line with 'STDOUT: Submitted batch job ID'.

    Parameters:
        command (str): The command to run for job submission.

    Returns:
        int: The extracted job ID.

    Raises:
        RuntimeError: If the command fails.
        ValueError: If the job ID cannot be found in the output.
    """
    try:
        result = subprocess.run(
            command, shell=True, capture_output=True, text=True, check=True
        )
    except subprocess.CalledProcessError as err:
        raise RuntimeError(
            f"Job submission failed: {err.stderr.strip()}"
        ) from err

    output = result.stdout.strip()
    # Extract job ID from a line like "Submitted job #12345"
    match = re.search(r"STDOUT: Submitted batch job (\d+)", output)
    if not match:
        raise ValueError(f"Could not extract job ID from output: {output}.")
    job_id = int(match.group(1))
    print(f"Submitted job {job_id}.")
    return job_id


def get_job_status(job_id: int) -> bool:
    """
    Runs sacct for job_id and checks if:
      - All jobs have known start and end times.
      - All jobs are done.
      - Every job's exit code is '0:0'.

    Returns:
        True if all jobs are done and have exit code '0:0'.
        False if not all jobs are finished (i.e. missing start or end times).

    Raises:
        ValueError: If all jobs are done but any job's exit code is not '0:0'.
    """
    # Run the sacct command and capture the output.
    command = f"sacct -j {job_id} -o JobID,State,ExitCode,Start,End --noheader"
    try:
        result = subprocess.run(
            command, shell=True, capture_output=True, text=True, check=True
        )
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"sacct command failed: {e.stderr.strip()}") from e

    output = result.stdout.strip()
    if not output:
        # If there's no output, nothing to process.
        return False

    # Process each line of output.
    lines = output.splitlines()
    all_done = True
    failing_jobs = (
        []
    )  # List to store (job_id, exit_code) for jobs with non-zero exit codes.

    for line in lines:
        # Split by whitespace; assumes each field (JobID, State, ExitCode, Start, End) is present.
        fields = line.split()
        if len(fields) < 5:
            continue  # Skip lines that do not have all required fields.
        id, state, exit_code, start, end = fields[:5]

        # Check if start and end times are known. You might need to adjust these
        # values depending on what sacct returns for unfinished jobs.
        if start == "Unknown" or end == "Unknown":
            all_done = False
            break

        # Check if the exit code is "0:0".
        if exit_code != "0:0" or state != "COMPLETED":
            failing_jobs.append((id, exit_code, state))

    # If not all jobs have finished, return False.
    if not all_done:
        return False

    # If all jobs are done but any job has a non-zero exit code, raise an error.
    if failing_jobs:
        for job, code, state in failing_jobs:
            print(
                f"Error: Job {job} failed with exit code {code} and state {state}.",
                file=sys.stderr,
            )
        raise ValueError(f"One or more jobs failed: {failing_jobs}")

    # All jobs are done and exited with 0:0.
    return True


def submit_and_monitor_job(
    commands: list[str],
    poll_interval: int = 10,
    retry_count: int = 10,
) -> None:
    """
    Submit a job using the provided command string and monitor it until it finishes.

    Parameters:
        command (str): The command to run for job submission.
        poll_interval (int): Time in seconds between status checks.
        retry_count (int): Number of times to retry monitoring the job, if any of the jobs fail.

    Raises:
        RuntimeError: If the job fails or ends with a state other than COMPLETED.
    """

    job_status = {command: (False, None) for command in commands}

    while True:
        for command, (done, job_id) in job_status.items():
            if not done and job_id is None:
                job_status[command] = (False, submit_job(command))
                print(f"Submitted job for {command}")

        for command, (done, job_id) in job_status.items():
            if not done:
                assert job_id is not None, "Job ID must be set."
                try:
                    new_done = get_job_status(job_id)
                    job_status[command] = (new_done, job_id)

                    if new_done:
                        print(f"Job {job_id} finished. Command: {command}")

                except ValueError as e:
                    print(f"ERROR: {e}", file=sys.stderr)
                    retry_count -= 1
                    print(f"Retries left: {retry_count}")
                    if retry_count <= 0:
                        raise RuntimeError(
                            "Job failed or ended with non-zero exit code."
                        )

                    # clear the job_id to resubmit
                    job_status[command] = (False, None)

        if all(done for done, _ in job_status.values()):
            print("All jobs finished.")
            return

        # polling delay
        time.sleep(poll_interval)


if __name__ == "__main__":
    args = parse_args()
    main(args)
