import logging
from math import prod

from typing import Callable, Optional
import uuid
import copy

from joblib import delayed
from tqdm import tqdm


from accelforge.frontend import arch
from accelforge.frontend._workload_isl._symbolic import get_stride_and_halo
from accelforge.frontend.renames import Rank, RankVariable
from accelforge.frontend.spec import Spec
from accelforge.frontend.mapping import Loop, Mapping, TensorHolder
from accelforge.util._frozenset import oset
from accelforge.frontend._workload_isl._isl import (
    get_rank_variable_bounds,
    get_tensor_size,
    get_operation_space_size,
)
from accelforge.frontend.workload import EinsumName, SymbolTable, TensorName, Workload

from accelforge.mapper.FFM._make_pmappings.make_pmapping_templates import (
    make_pmapping_templates,
)
from accelforge.frontend.mapper.metrics import Metrics
from accelforge.mapper.FFM._make_pmappings.make_pmappings_from_templates import (
    make_pmappings_from_templates,
)
from accelforge.mapper.FFM._join_pmappings.compatibility import Compatibility
from accelforge.mapper.FFM._join_pmappings.pmapping_group import PmappingGroup
from accelforge.mapper.FFM._make_pmappings.make_pmappings_from_templates.symbol_relations import (
    get_initial_delta_choices,
)
from accelforge.util.parallel import (
    parallel,
    _memmap_read,
    get_n_parallel_jobs,
    is_using_parallel_processing,
)
from accelforge.mapper.FFM._make_pmappings.pmapper_job import (
    Job,
    SameCompatibilityJobs,
)


def get_rank_variable_bounds_for_all_einsums(spec: Spec):
    rank_variable_bounds = {
        einsum_name: get_rank_variable_bounds(spec.workload, einsum_name)
        for einsum_name in spec.workload.einsum_names
    }
    result = {}
    for e1, rv1 in rank_variable_bounds.items():
        result.update(rv1)
        for e2, rv2 in rank_variable_bounds.items():
            for r in oset(rv1.keys()) & oset(rv2.keys()):
                if rv1[r] != rv2[r]:
                    raise ValueError(
                        f"Rank variable {r} has different bounds for "
                        f"einsum {e1} and {e2}: {rv1[r]} and {rv2[r]}"
                    )
    return result


def get_n_computes(spec: Spec, einsum_name: EinsumName | None = None) -> int:
    einsums = spec.workload.einsums
    einsums = [einsum_name] if einsum_name is not None else spec.workload.einsum_names
    return sum(get_operation_space_size(spec.workload, e) for e in einsums)


def get_per_tensor_size(
    workload: Workload, return_n_elements: bool = False
) -> dict[TensorName, int]:
    sizes = {}
    for einsum in workload.einsums:
        for tensor in einsum.tensor_names:
            sizes[tensor] = get_tensor_size(workload, tensor)
            if not return_n_elements:
                sizes[tensor] *= einsum.tensor_accesses[tensor].bits_per_value
    return sizes


def get_jobs(
    spec: Spec,
    metrics: Metrics,
    einsum_names: list[EinsumName],
    fail_if_no_pmappings_for_einsum: bool,
    print_progress: bool,
) -> dict[EinsumName, dict[Compatibility, SameCompatibilityJobs]]:

    spec = spec

    einsum2jobs = {}
    fusable_tensors = spec.workload.tensor_names_used_in_multiple_einsums
    rank_variable_bounds = get_rank_variable_bounds_for_all_einsums(spec)

    einsum2spec: dict[EinsumName, Spec] = {}
    s = "Getting energy, latency, and leak power for components running each Einsum. "

    def _get_per_einsum_spec(spec, einsum_name):
        result = (
            spec._spec_eval_expressions(
                einsum_name=einsum_name,
                eval_arch=True,
                eval_non_arch=False,
            )
            .calculate_component_area_energy_latency_leak(
                einsum_name=einsum_name,
                area=False,
            )
            ._for_einsum(einsum_name)
            ._clear_component_models()
        )
        return einsum_name, _memmap_read(result)

    for einsum_name, result in parallel(
        [
            delayed(_get_per_einsum_spec)(spec, einsum_name)
            for einsum_name in einsum_names
        ],
        pbar=s if print_progress else None,
    ):
        einsum2spec[einsum_name] = result

    def make_jobs_for_einsum(
        einsum_name: EinsumName,
        spec: Spec,
        stride_and_halo: dict[
            EinsumName,
            dict[TensorName, dict[tuple[Rank, RankVariable], tuple[int, int]]],
        ],
        initial_delta_choices_for_einsum: dict[RankVariable, frozenset[int]],
        n_einsums: int,
        intermediate_tensors: set[TensorName],
    ):
        jobs = {}
        workload_einsum = spec.workload.einsums[einsum_name]
        for flattened_arch in spec._get_flattened_architecture(einsum_name=einsum_name):
            # Create jobs for each Einsum
            job = Job(
                spec_one_einsum=spec,
                einsum_name=einsum_name,
                metrics=metrics,
                rank_variable_bounds=rank_variable_bounds,
                flattened_arch=_memmap_read(flattened_arch),
                job_id=uuid.uuid4(),
                fusable_tensors=fusable_tensors & workload_einsum.tensor_names,
                stride_and_halo=stride_and_halo,
                initial_delta_choices=initial_delta_choices_for_einsum,
                objective_tolerance=spec.mapper.objective_tolerance,
                resource_usage_tolerance=spec.mapper.resource_usage_tolerance,
                workload_n_einsums=n_einsums,
                intermediate_tensors=intermediate_tensors
                & workload_einsum.tensor_names,
                explore_imperfect_spatial_loops=bool(
                    spec.mapper.explore_imperfect_spatial_loops
                ),
                explore_imperfect_temporal_loops=bool(
                    spec.mapper.explore_imperfect_temporal_loops
                ),
            )
            for j in make_pmapping_templates(job, print_progress):
                jobs.setdefault(j.compatibility, SameCompatibilityJobs()).append(j)

        return einsum_name, jobs

    stride_and_halo = get_stride_and_halo(spec.workload)
    initial_delta_chocies = {
        e: get_initial_delta_choices(e, spec.workload)
        for e in spec.workload.einsum_names
    }

    input_tensors = oset.union(*[e.input_tensor_names for e in spec.workload.einsums])
    output_tensors = oset.union(*[e.output_tensor_names for e in spec.workload.einsums])
    intermediate_tensors = input_tensors & output_tensors

    n_einsums = len(spec.workload.einsum_names)
    for einsum_name, jobs in parallel(
        [
            delayed(make_jobs_for_einsum)(
                einsum_name,
                spec._for_einsum(einsum_name),
                stride_and_halo,
                initial_delta_chocies[einsum_name],
                n_einsums,
                intermediate_tensors,
            )
            for einsum_name, spec in einsum2spec.items()
        ],
        pbar="Generating jobs" if print_progress else None,
    ):
        einsum2jobs.setdefault(einsum_name, {})
        for compatibility, job_list in jobs.items():
            einsum2jobs[einsum_name].setdefault(
                compatibility, SameCompatibilityJobs()
            ).extend(job_list)

    if fail_if_no_pmappings_for_einsum:
        for einsum_name, jobs in einsum2jobs.items():
            if len(jobs) == 0:
                raise ValueError(
                    f"No pmappings for {einsum_name}. Was the mapspace overconstrained?"
                )

    total_jobs = sum(len(jobs) for jobs in einsum2jobs.values())
    n_procs = get_n_parallel_jobs()
    memory_limit = min(
        spec.mapper.memory_limit, spec.mapper.memory_limit_per_process / n_procs
    )
    time_limit = min(
        spec.mapper.time_limit * n_procs / max(total_jobs, 1),
        spec.mapper.time_limit_per_pmapping_template,
    )
    total_templates = 0
    for einsum_name, compatibility_jobs in einsum2jobs.items():
        n_compatibility_jobs = sum(len(j) for j in compatibility_jobs.values())
        total_templates += n_compatibility_jobs
        if print_progress:
            s = "jobs" if n_compatibility_jobs > 1 else "job"
            print(f"Einsum {einsum_name} has {n_compatibility_jobs} pmapping {s}:")
        i = 0
        for job_list in compatibility_jobs.values():
            for job in job_list:
                if print_progress:
                    print(f"\t{i}\t{job.mapping.compact_str()}")
                    i += 1
                job.memory_limit = memory_limit
                job.time_limit = time_limit

    if print_progress:
        s = "templates" if total_templates > 1 else "template"
        print(f"Total number of pmapping {s}: {total_templates}")

    return einsum2jobs


def get_memories_to_track(
    spec: Spec,
    einsum2jobs: dict[EinsumName, list[Job]],
    metrics: Metrics,
    can_combine_multiple_runs: bool,
    print_progress: bool = True,
) -> tuple[list[str], list[str]]:

    memories_track_all = oset()
    for einsum, jobs in einsum2jobs.items():
        for job in jobs:
            memories_track_all.update(
                m.name for m in job.flattened_arch if isinstance(m, arch.Memory)
            )

    memories_track_pmappings_only = []
    ignored_resources = oset()

    # If we're combining the pmappings from multiple runs, we can't conclude anything
    # about the metrics to track
    if can_combine_multiple_runs or metrics.RESOURCE_USAGE in metrics:
        ignored_resources = oset()
        return (
            memories_track_all,
            memories_track_pmappings_only,
            ignored_resources,
        )

    tensor_sizes = {}
    for tensor, size in get_per_tensor_size(spec.workload).items():
        scale = 1
        for einsum in spec.workload.einsums_with_tensor(tensor):
            if einsum.tensor_accesses[tensor].persistent:
                scale = max(scale, spec.workload.n_instances * einsum.n_instances)
        tensor_sizes[tensor] = size * scale

    # If the memory is big enough to hold all the tensors then we don't need to consider
    # it
    for memory in list(memories_track_all):
        usage = 0
        for einsum in einsum2jobs.keys():
            job = einsum2jobs[einsum][0]
            try:
                mem: arch.Memory = job.spec_one_einsum.arch.find(memory)
            except ValueError:
                continue
            for tensor in spec.workload.einsums[einsum].tensor_names:
                if mem.size == 0:
                    usage = 2  # FAIL
                else:
                    workload_bpv = (
                        spec.workload.einsums[einsum]
                        .tensor_accesses[tensor]
                        .bits_per_value
                    )
                    effective_bpv = mem.bits_per_value.get(tensor, workload_bpv)
                    usage += (
                        tensor_sizes[tensor] * effective_bpv / (workload_bpv * mem.size)
                    )

        if usage <= 1:
            ignored_resources.add(memory)
            if print_progress:
                print(
                    f"Not tracking memory {memory}. It is big enough to simultaneously "
                    f"hold every workload tensor that may be stored in it. Max "
                    f"possible usage: {usage * 100:.2f}%"
                )
            memories_track_all.remove(memory)

    # If the memory is below every backing tensor holder node, then we need it for the
    # pmapping exploration but can drop it immediately
    for m in list(memories_track_all):
        must_track = False
        for _, jobs in einsum2jobs.items():
            for job in jobs:
                seen = False
                for node in job.mapping.nodes:
                    if isinstance(node, TensorHolder) and node.component == m:
                        seen = True
                        if node.persistent:
                            ignored_resources.add(m)
                        if node._backing:
                            must_track = True
                    if isinstance(node, Loop) and node._fused and seen:
                        must_track = True

        if not must_track:
            memories_track_all.remove(m)
            memories_track_pmappings_only.append(m)
            if print_progress:
                print(
                    f"Not tracking memory {m} across joining stages. It is never "
                    f"reserved across fused loop iterations."
                )

    return memories_track_all, memories_track_pmappings_only, ignored_resources


def make_pmappings(
    spec: Spec,
    can_combine_multiple_runs: bool,
    metrics: Metrics = Metrics.ENERGY | Metrics.LATENCY,
    einsum_names: Optional[list[EinsumName]] = None,
    fail_if_no_pmappings_for_einsum: bool | None = None,
    print_progress: bool = True,
    one_pbar_only: bool = False,
) -> tuple[
    dict[EinsumName, list[PmappingGroup]],
    dict[EinsumName, dict[uuid.UUID, Mapping]],
    dict[EinsumName, list[Job]],
]:
    """
    Explores pmapspace of `einsum_names` (default: all Einsums in workload).
    """
    spec = copy.deepcopy(spec)

    if einsum_names is None:
        einsum_names = spec.workload.einsum_names

    if one_pbar_only:
        assert (
            not print_progress
        ), "print_progress must be False if one_pbar_only is True"

    if fail_if_no_pmappings_for_einsum is None:
        fail_if_no_pmappings_for_einsum = not can_combine_multiple_runs

    spec = spec._spec_eval_expressions(
        eval_arch=False,
        eval_non_arch=True,
    )

    einsum2jobs = {}
    new_einsum2jobs = get_jobs(
        spec,
        metrics,
        einsum_names,
        fail_if_no_pmappings_for_einsum,
        print_progress,
    )
    _fill_jobs_with_memories_to_track(
        new_einsum2jobs, spec, metrics, can_combine_multiple_runs, print_progress
    )
    for einsum_name, jobs in new_einsum2jobs.items():
        einsum2jobs.setdefault(einsum_name, {})
        for compatibility, job_list in jobs.items():
            einsum2jobs[einsum_name].setdefault(
                compatibility, SameCompatibilityJobs()
            ).extend(job_list)

    calls = _allocate_jobs(einsum2jobs, print_progress)

    # Sort the calls by the length of the longest mapping in each job. We get long
    # poles with the long mappings, so we want to get them done early so we don't
    # have one or two procs slowing us down at the end.
    def get_longest_mapping_length(call):
        j: SameCompatibilityJobs = call[2]["jobs_with_similar_compatibilities"]
        return max([len(j2.mapping.nodes) for j2 in j])

    calls = sorted(calls, key=get_longest_mapping_length, reverse=True)

    pmapping_objects = {}
    pmapping_groups = {einsum_name: [] for einsum_name in spec.workload.einsum_names}
    keep_rates = []
    for (
        einsum_name,
        new_pmapping_groups,
        pmappings,
        cur_keep_rates,
    ) in parallel(
        calls,
        pbar=f"Generating pmappings" if print_progress or one_pbar_only else None,
        return_as="generator_unordered",
    ):
        pmapping_groups[einsum_name].extend(new_pmapping_groups)
        pmapping_objects.setdefault(einsum_name, {}).update(pmappings)
        keep_rates.extend(cur_keep_rates)

    for einsum_name in list(pmapping_groups.keys()):
        pmapping_groups[einsum_name] = PmappingGroup.combine_combineable(
            pmapping_groups[einsum_name],
            "All",
            pbar_postfix=f" for {einsum_name}",
            print_progress=print_progress,
        )

    return_jobs = {
        e: [j for jobs in v.values() for j in jobs] for e, v in einsum2jobs.items()
    }

    # Propagate keep_rates from worker-side job copies back to original jobs
    id2keep_rates = {jid: (keep_rates, n) for jid, keep_rates, n in keep_rates}
    for jobs in return_jobs.values():
        for job in jobs:
            assert job.job_id in id2keep_rates, f"Job {job.job_id} not found"
            job.pmapping_keep_rates.update(id2keep_rates[job.job_id][0])
            job.n_total_pmappings = id2keep_rates[job.job_id][1]

    return pmapping_groups, pmapping_objects, return_jobs


def _raise_error_if_no_pmappings(einsum2jobs):
    for einsum_name, jobs in einsum2jobs.items():
        if len(jobs) == 0:
            raise ValueError(
                f"No pmappings for {einsum_name}. " f"Was the mapspace overconstrained?"
            )


def _allocate_jobs(einsum2jobs, print_progress: bool = True):
    calls = []
    for einsum_name, jobs in einsum2jobs.items():
        calls.extend(
            delayed(make_pmappings_from_templates)(
                jobs_with_similar_compatibilities=job_list,
            )
            for job_list in jobs.values()
        )

    split = True

    if split:
        calls = []
        for einsum_name, jobs in einsum2jobs.items():
            for job_list in jobs.values():
                calls.extend(
                    delayed(make_pmappings_from_templates)(
                        jobs_with_similar_compatibilities=job,
                    )
                    for job in job_list.split()
                )
    return calls


def _fill_jobs_with_memories_to_track(
    einsum2jobs: dict[EinsumName, dict[Compatibility, SameCompatibilityJobs]],
    spec,
    metrics,
    can_combine_multiple_runs,
    print_progress,
):
    einsum2jobs_flattened = {
        e: [j for jobs in v.values() for j in jobs] for e, v in einsum2jobs.items()
    }

    memories_track_all, memories_track_pmappings_only, ignored_resources = (
        get_memories_to_track(
            spec,
            einsum2jobs_flattened,
            metrics,
            can_combine_multiple_runs,
            print_progress,
        )
    )
    for jobs in einsum2jobs_flattened.values():
        for j in jobs:
            j.memories_track_all = memories_track_all
            j.memories_track_pmappings_only = memories_track_pmappings_only
            j.ignored_resources = ignored_resources
