import functools
from math import prod
import time

from accelforge._accelerated_imports import pandas as pd

from paretoset import paretoset
from joblib import delayed
from sympy import factorint

from accelforge._accelerated_imports import np
from accelforge.util._frozenset import oset
from accelforge.util.parallel import parallel

from accelforge.mapper.FFM._pareto_df.df_convention import (
    col2reservation,
    col_used_in_pareto,
    is_fused_loop_col,
    is_n_iterations_col,
    is_objective_col,
)

from .fast_pareto import fast_pareto_mask


def dominates(a: pd.Series, b: pd.Series) -> bool:
    return all(a[i] <= b[i] for i in range(len(a)))


def check_dominance(df: pd.DataFrame, n_optimal: int):
    # mask = np.zeros(len(df), dtype=bool)
    # mask[:new_point] = True
    mask = np.zeros(len(df) - n_optimal, dtype=bool)
    for col in df.columns:
        compare = df.iloc[n_optimal - 1][col]
        mask = mask | (df[col].iloc[n_optimal:] < compare)
    return np.concatenate([np.ones(n_optimal, dtype=bool), mask])


def quickpareto(df: pd.DataFrame) -> pd.DataFrame:
    # Step 1: Sort by the column with the most unique values
    # Step 2: Extract the first row. Add it to the pareto set
    # Step 3: Remove all dominated points
    # Step 4: Repeat until no more points to add

    # Step 1: Sort by the column with the most unique values
    original_len = len(df)
    col_to_sort = max(df.columns, key=lambda c: df[c].nunique())
    df = df.sort_values(by=col_to_sort).drop(columns=[col_to_sort])

    new_point = 0
    while new_point < len(df):
        mask = check_dominance(df, new_point + 1)
        df = df[mask]
        new_point += 1

    # Turn the index into a mask
    mask = np.zeros(original_len, dtype=bool)
    mask[df.index] = True
    return mask


def makepareto_quick2(mappings: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
    from fast_pareto import is_pareto_front

    m2 = mappings[columns]
    m2 = m2[is_pareto_front(m2.to_numpy())].drop_duplicates()
    return mappings.loc[m2.index]


def makepareto_quick(mappings: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
    return mappings[quickpareto(mappings[columns])]


def paretofy_chunk(chunk, sense: list[str]):
    return fast_pareto_mask(chunk.values if hasattr(chunk, "values") else chunk, sense)


def makepareto_merge(
    mappings: pd.DataFrame,
    columns: list[str],
    parallelize: bool = False,
    split_by_cols: list[str] = (),
) -> pd.DataFrame:
    if len(mappings) <= 1:
        return mappings

    sense = ["min"] * len(columns) + ["diff"] * len(split_by_cols)
    combined = mappings[columns + list(split_by_cols)]
    return mappings[fast_pareto_mask(combined.values, sense)]


def makepareto_time_compare(mappings: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
    t0 = time.time()
    pareto = makepareto_merge(mappings, columns)
    t1 = time.time()
    merge_time = t1 - t0
    print(
        f"Time to make pareto with merge: {t1 - t0: .2f}. Number of pareto points: {len(pareto)}"
    )

    t0 = time.time()
    pareto2 = makepareto_quick2(mappings, columns)
    t1 = time.time()
    print(
        f"Time to make pareto with quick: {t1 - t0: .2f}. Number of pareto points: {len(pareto2)}"
    )
    quick_time = t1 - t0

    print(f"Quick is {quick_time / merge_time: .2f}x slower")

    if len(pareto) != len(pareto2):
        print(f"mismatch: {len(pareto)} != {len(pareto2)}")
        makepareto_quick2(mappings)

    return pareto2


# 2d. Blockwise vectorized CuPy Pareto front with sorting by one objective (full check)
# 2c. Fully vectorized CuPy brute-force Pareto front
# (returns numpy mask for compatibility)
def pareto_front_cupy_vectorized(X):
    # if len(X) > 1000:
    #     return X[paretoset(X.get(), sense=["min"] * X.shape[1])]

    # Broadcast X_gpu to (n, n, m) for all-pairs comparison
    A = X[:, None, :]  # shape (n, 1, m)
    B = X[None, :, :]  # shape (1, n, m)
    less_equal = (B <= A).all(axis=2)  # shape (n, n)
    strictly_less = (B < A).any(axis=2)  # shape (n, n)
    dominated = less_equal & strictly_less  # shape (n, n)
    is_pareto = ~dominated.any(axis=1)
    return is_pareto


# 2d. Recursive blockwise merge CuPy Pareto front with sorting by one objective
def pareto_front_cupy_blockwise_sorted_recursive(X, block_size=2000):
    N = X.shape[0]
    if N <= block_size:
        # Base case: just compute Pareto front directly
        mask = pareto_front_cupy_vectorized(X)
        return mask
    # Split into two halves
    mid = N // 2
    a, b = X[:mid], X[mid:]
    mask_a = pareto_front_cupy_blockwise_sorted_recursive(a, block_size)
    mask_b = pareto_front_cupy_blockwise_sorted_recursive(b, block_size)
    # Get Pareto-optimal points from both halves
    pareto_points_a = a[mask_a]
    pareto_points_b = b[mask_b]
    merged_points = np.vstack([pareto_points_a, pareto_points_b])
    # Compute Pareto front of the merged set
    merged_mask = pareto_front_cupy_vectorized(merged_points)
    merged_indices = np.where(merged_mask)[0]
    # Map merged_indices back to the original indices in X
    # First, get the indices in X for the merged points
    indices_a = np.where(mask_a)[0]
    indices_b = np.where(mask_b)[0] + mid
    all_indices = np.concatenate([indices_a, indices_b])
    merged_indices_in_X = all_indices[merged_indices]
    # Build the final mask for X
    mask = np.zeros(N, dtype=bool)
    mask[merged_indices_in_X] = True
    return mask


# def makepareto(
#     mappings: pd.DataFrame,
#     columns: list[str] = None,
#     parallelize: bool = False,
#     split_by_cols: list[str] = (),
# ) -> pd.DataFrame:
#     # return makepareto_time_compare(mappings)
#     if columns is None:
#         columns = [c for c in mappings.columns if col_used_in_pareto(c)]
#     if _accelerated_imports.ACCELERATED:
#         mask = pareto_front_cupy_blockwise_sorted_recursive(mappings[columns].to_cupy())
#         return mappings[mask]


def round_to_tolerance(x: pd.Series, tolerance: float) -> pd.Series:
    if tolerance == 0:
        return x
    x = np.round(x / tolerance) * tolerance
    return x


def logscale_to_tolerance(x: pd.Series, tolerance: float) -> pd.Series:
    if tolerance == 0 or tolerance is None:
        return x
    assert tolerance > 0
    if x.min() <= 0:
        return x
    rounded = np.round(np.log(x) / np.log(1 + tolerance))
    return np.exp(rounded * np.log(1 + tolerance))


def multi_round(
    x: np.ndarray, tolerance: float, absolute_tolerance: float
) -> np.ndarray:
    tolerance = 0 if tolerance is None else tolerance
    absolute_tolerance = 0 if absolute_tolerance is None else absolute_tolerance
    if tolerance == 0:
        return round_to_tolerance(x, absolute_tolerance)
    if absolute_tolerance == 0:
        return logscale_to_tolerance(x, tolerance)

    abs_step_size = absolute_tolerance
    log_step_size = tolerance * x

    rounded = round_to_tolerance(x, abs_step_size)
    logscale_rounded = logscale_to_tolerance(x, tolerance)

    use_log_mask = log_step_size > abs_step_size
    new = np.where(use_log_mask, logscale_rounded, rounded)
    return new


def makepareto(
    mappings: pd.DataFrame,
    columns: list[str] = None,
    split_by_cols: list[str] = (),
    resource_usage_tolerance: float = 0,
    objective_tolerance: float = 0,
    absolute_resource_usage_tolerance: float = 0,
) -> pd.DataFrame:
    # return makepareto_time_compare(mappings)
    if columns is None:
        columns = [c for c in mappings.columns if col_used_in_pareto(c)]

    # Number of iterations is derived from the tile shapes, so we don't need to use it,
    # since any row with the same tile shapes will have the same number of iterations.
    split_by_cols = list(split_by_cols) + [
        c
        for c in mappings.columns
        if is_fused_loop_col(c) and not is_n_iterations_col(c)
    ]

    goals = []
    to_pareto = []
    for c in mappings.columns:
        if mappings[c].nunique() <= 1:
            continue

        if c in columns and is_objective_col(c):
            to_pareto.append(logscale_to_tolerance(mappings[c], objective_tolerance))
            goals.append("min")
        elif c in split_by_cols:
            to_pareto.append(mappings[c])
            goals.append("diff")
        elif c in columns:
            x = mappings[c]
            if col2reservation(c) is not None:
                x = multi_round(
                    x, resource_usage_tolerance, absolute_resource_usage_tolerance
                )
            to_pareto.append(x)
            goals.append("min")

    if not to_pareto:
        return mappings.iloc[0:1]

    # import os
    # import uuid
    # import pickle as pkl

    # os.makedirs("/tmp/paretos", exist_ok=True)
    # pkl.dump(
    #     (pd.concat(to_pareto, axis=1), goals),
    #     open(f"/tmp/paretos/{uuid.uuid4()}.pkl", "wb")
    # )

    combined = pd.concat(
        (pd.Series(p).reset_index(drop=True) for p in to_pareto), axis=1
    )
    return mappings[fast_pareto_mask(combined.values, goals)]


@functools.lru_cache(maxsize=10000)
def _factorint_cached(x: int):
    return factorint(x)


def prime_factor_counts(arr: np.ndarray) -> np.ndarray:
    arr = np.asarray(arr, dtype=int)
    unique_vals = np.unique(arr)
    factorizations = {x: _factorint_cached(x) for x in unique_vals}

    # Gather all unique primes
    all_primes = sorted(oset(p for f in factorizations.values() for p in f))

    # Build result matrix
    result = np.zeros((len(arr), len(all_primes)), dtype=int)
    prime_index = {p: j for j, p in enumerate(all_primes)}

    for i, x in enumerate(arr.flatten()):
        for p, exp in factorizations[x].items():
            result[i, prime_index[p]] = exp

    return result


def makepareto_numpy(
    mappings: np.ndarray,
    goals: list[str],
    dirty: bool = False,  # Doesn't do anything
    tolerances: list[float] = None,
    absolute_tolerances: list[float] = None,
) -> pd.DataFrame:

    to_pareto = []
    new_goals = []
    if tolerances is None:
        tolerances = [0] * len(goals)
    if absolute_tolerances is None:
        absolute_tolerances = [0] * len(goals)

    assert len(goals) == mappings.shape[1]
    for c in range(mappings.shape[1]):
        if len(np.unique(mappings[:, c])) <= 1:
            continue

        goal = goals[c]
        rounded = multi_round(mappings[:, c], tolerances[c], absolute_tolerances[c])

        if goal in ["min", "max"]:
            to_pareto.append(rounded if goal == "min" else -rounded)
            new_goals.append("min")
        elif goal == "diff":
            to_pareto.append(rounded)
            new_goals.append("diff")
        elif goal == "min_per_prime_factor":
            counts = prime_factor_counts(rounded)
            for i in range(counts.shape[1]):
                to_pareto.append(counts[:, i])
                new_goals.append("min")
        elif goal == "max_per_prime_factor":
            counts = prime_factor_counts(rounded)
            for i in range(counts.shape[1]):
                to_pareto.append(counts[:, i])
                new_goals.append("max")
        else:
            raise ValueError(f"Unknown goal: {goal}")

    if not to_pareto:
        n = np.zeros(mappings.shape[0], dtype=bool)
        n[0] = True
        return n

    data = np.concatenate([x.reshape(-1, 1) for x in to_pareto], axis=1)
    return fast_pareto_mask(data, new_goals)
