import numpy as np
import cvxpy as cp
import time
from config import *  # Import all parameters from config.py
class ResourceAllocator:
    """
        Given a fixed unloading matrix X,
        The solver finds the optimal (a, tau, e, f, T_ex) to minimize the weighted sum.
    """

    def __init__(self):
        """Initialize the solver"""
        # Estimate some constants
        self.C_rate = (np.log(2) * V_U) / BANDWIDTH_HZ
        print(f"Convex ResourceAllocator (CVXPY) initialized for {N_WDs} WDs.")

    def _calculate_loads(self, action_matrix: np.ndarray, task_sizes: np.ndarray) -> np.ndarray:
        """
       Calculate the total computing load(L_j) of each WD j 。
        L_j = (local tasks of WD j) + (All other tasks offloaded by WD i to WD j)
        """
        N = N_WDs
        loads = np.zeros(N)
        for j in range(N):  # Target WD j (node index j+1)
            for i in range(N):  # Source: WD i (Task: c_i)
                # action_matrix[i, j+1] is x_{i+1, j+1}
                loads[j] += task_sizes[i] * action_matrix[i, j + 1]
        return loads

    def solve(self, action_matrix: np.ndarray, channel_matrix: np.ndarray, task_sizes: np.ndarray) -> dict:
        """
       Solve the convex subproblem using CVXPY

        Args:
            action_matrix (np.ndarray): Offload decision X, shape (N_WDs, N_NODES)
            channel_matrix (np.ndarray): Channel gain H, shape (N_NODES, N_NODES)
            task_sizes (np.ndarray): Task size c, shape (N_WDs,)

        Returns:
            dict: Solution result
        """
        N = N_WDs
        N_nodes = N_NODES
        start_opt_time = time.time()

        try:
            # --- 1. Pre-computed constant ---
            # L_j: Each WD (j=0...) N-1, corresponding to node 1.. The total computing load (bits) of N)
            L_j = self._calculate_loads(action_matrix, task_sizes)

            # s_ij: The number of bits sent from WD i (node i+1) to node j
            s_ij = np.zeros((N, N_nodes))
            for i in range(N):
                for j in range(N_nodes):
                    if i + 1 == j:  # Skip local computing
                        continue
                    s_ij[i, j] = task_sizes[i] * action_matrix[i, j]

            # h_ij: Channel gain from WD i (node i+1) to node j
            h_ij = channel_matrix[1:, :]

            # gamma_ij = h_ij / N0
            gamma_ij = h_ij / NOISE_POWER

            # E_harv_j (coefficient): the energy collected by WD j (node j+1) from AP (node 0)
            # E_harv_j = (coefficient) * a
            E_harv_coeff_j = (ENERGY_HARVEST_EFFICIENCY * AP_TX_POWER_W *
                              ENERGY_HARVEST_SCALE_S * channel_matrix[0, 1:])

            # --- 2. Define the CVXPY variable ---
            # Scalar
            a = cp.Variable(pos=True, name="a")
            T_ex = cp.Variable(pos=True, name="T_ex")
            # Vector (N_WDs)
            f = cp.Variable(N, pos=True, name="f")
            # Matrix (N_WDs, N_NODES)
            tau = cp.Variable((N, N_nodes), pos=True, name="tau")
            e = cp.Variable((N, N_nodes), pos=True, name="e")



            # --- 3. Define the objective function---
            # T_total = a + T_offload + T_execute
            T_offload = cp.sum(tau)
            T_total = a + T_offload + T_ex

            # E_total = E_tx + E_comp
            E_tx = cp.sum(e)
            E_comp = cp.sum(CYCLES_PER_BIT * ENERGY_EFFICIENCY_COEFF * cp.multiply(L_j, cp.power(f, 2)))
            E_total = E_tx + E_comp

            objective = cp.Minimize(TIME_WEIGHT * T_total + ENERGY_WEIGHT * E_total)

            # --- 4. Define constraints ---
            constraints = []

            # Variable boundary
            TAU_FLOOR = 1e-9
            constraints += [a >= TAU_FLOOR, a <= 1.0]
            constraints += [f >= CPU_FREQ_MIN_HZ, f <= CPU_FREQ_MAX_HZ]


            # Power limit(P_min <= e/tau <= P_max)
            constraints += [e >= TX_POWER_MIN_W * tau]
            constraints += [e <= TX_POWER_MAX_W * tau]

            # Execution time constraint (T_ex >= L_j * phi / f_j)
            constraints += [T_ex >= cp.multiply(L_j * CYCLES_PER_BIT, cp.inv_pos(f))]

            # Energy neutrality constraint (E_self_j <= E_harv_j)
            E_tx_j = cp.sum(e, axis=1)  # The total transmission energy consumption of WD j (i=0..N-1) 
            E_comp_j = CYCLES_PER_BIT * ENERGY_EFFICIENCY_COEFF * cp.multiply(L_j, cp.power(f, 2))
            E_self_j = E_tx_j + E_comp_j
            E_harv_j = E_harv_coeff_j * a
            constraints += [E_self_j <= E_harv_j]

            # Rate constraint (s_ij <= C * tau * log2(1 + gamma * e / tau))
            #Use -rel_entr(tau, tau + gamma*e) to represent tau * log((tau + gamma*e)/tau)
            for i in range(N):
                for j in range(N_nodes):
                    if i + 1 == j:
                        # The constraint self-transport is 0
                        constraints += [tau[i, j] == 0, e[i, j] == 0]
                        continue

                    if s_ij[i, j] < EPS:
                        # If there is no data to send, no time and energy will be allocated
                        constraints += [tau[i, j] == 0, e[i, j] == 0]
                        continue

                    constraints += [tau[i, j] >= TAU_FLOOR]
                    # LHS = s_ij * log(2) * v_u / B
                    LHS = s_ij[i, j] * self.C_rate
                    # RHS = tau * log(1 + gamma * e / tau)
                    RHS = -cp.rel_entr(tau[i, j], tau[i, j] + gamma_ij[i, j] * e[i, j])
                    constraints += [LHS <= RHS]

            # --- 5. Solve the problem ---
            problem = cp.Problem(objective, constraints)
            # Use ECOS or SCS, which are designed for this type of problem
            #problem.solve(solver=cp.ECOS, verbose=False)
            problem.solve(solver=cp.SCS, eps=1e-5, max_iters=20000, warm_start=True, verbose=True)
            print("status=", problem.status, "solver=", problem.solver_stats.solver_name)

            opt_duration = time.time() - start_opt_time

            if problem.status in ["optimal", "optimal_inaccurate"]:
                obj_val = problem.value
                t_total_val = (a.value + np.sum(tau.value) + T_ex.value)
                e_total_val = (
                            np.sum(e.value) + np.sum(CYCLES_PER_BIT * ENERGY_EFFICIENCY_COEFF * L_j * (f.value ** 2)))

                metrics = {
                    "success": True,
                    "message": problem.status,
                    "objective": obj_val,
                    "total_time": t_total_val,
                    "total_energy": e_total_val,
                    "a": a.value,
                    "cpu_freqs": f.value,
                    "T_harvest": a.value,
                    "T_offload": np.sum(tau.value),
                    "T_execute": T_ex.value,
                    "opt_duration_s": opt_duration
                }
            else:
                # Solution failed
                metrics = {
                    "success": False,
                    "message": problem.status,
                    "objective": np.inf,
                    "opt_duration_s": opt_duration
                }

        except Exception as e:
            # The CVXPY solver may fail completely due to numerical issues
            print(f"Warning: CVXPY solver failed with exception: {e}")
            metrics = {
                "success": False,
                "message": str(e),
                "objective": np.inf,
                "opt_duration_s": time.time() - start_opt_time
            }

        return metrics


# Example usage (within the file for direct testing)
if __name__ == '__main__':
    # test
    chan_matrix = np.array([[0.00000000e+00, 6.62546102e-06, 2.31014624e-06, 1.83911263e-06],
                            [6.62546102e-06, 0.00000000e+00, 4.80614196e-06, 1.01083395e-06],
                            [2.31014624e-06, 4.80614196e-06, 0.00000000e+00, 3.81168776e-06],
                            [1.83911263e-06, 1.01083395e-06, 3.81168776e-06, 0.00000000e+00]])
    tasks = np.array([46170.08926378, 97076.77384453, 12340.47028163])
    dummy_action = np.array([[0.17732602, 0.11025, 0.28628772, 0.42613626],
                             [0.358793, 0.21895458, 0.21797714, 0.20427531],
                             [0.10907146, 0.13038586, 0.6538224, 0.10672026]])

    print("\n--- Testing Resource Allocator (CVXPY Convex Solver) ---")
    print(f"Target Dummy Action:\n{dummy_action}")
    print(f"Task sizes (Mbits): {tasks / 1e6}")

    # Instantiate and solve
    allocator = ResourceAllocator()
    results = allocator.solve(dummy_action, chan_matrix, tasks)

    print("\n--- Optimization Results ---")
    if results["success"]:
        print(f"Success: {results['success']} ({results['message']})")
        print(f"Optimization duration: {results['opt_duration_s']:.4f} s")
        print(f"Optimal Objective: {results['objective']:.6f}")
        print(f"  Total Time: {results['total_time']:.6f}")
        print(f"  Total Energy: {results['total_energy']:.6e}")
        print(f"Optimal 'a': {results['a']:.6f}")
        print(f"Optimal 'f' (MHz): {results['cpu_freqs'] / 1e6}")
        print(f"  T_harvest: {results['T_harvest']:.6f}")
        print(f"  T_offload: {results['T_offload']:.6f}")
        print(f"  T_execute: {results['T_execute']:.6f}")
    else:
        print(f"Optimization Failed: {results['message']}")

# -----------------
# System Parameters
# -----------------
N_WDs = 3  # Number of Wireless Devices (WDs)
N_NODES = N_WDs + 1  # Total number of nodes (AP + WDs)

# --- Task Model ---
TASK_SIZE_MIN_KB = 50  # Minimum task size in KiloBytes
TASK_SIZE_MAX_KB = 100 # Maximum task size in KiloBytes
TASK_SIZE_MIN = TASK_SIZE_MIN_KB * 1000  # Minimum task size in bits
TASK_SIZE_MAX = TASK_SIZE_MAX_KB * 1000  # Maximum task size in bits
CYCLES_PER_BIT = 10 # Number of CPU cycles required to process one bit (phi)


NOISE_POWER = 1e-10    # Noise power in Watts (N0)
BANDWIDTH_HZ = 3e6     # Communication bandwidth in Hz (B)
V_U = 1.1              # Communication overhead factor (v_u)

# --- Energy Model ---
ENERGY_HARVEST_EFFICIENCY = 0.8  # Efficiency of energy harvesting (mu)
AP_TX_POWER_W = 6.0            # AP's transmit power for energy harvesting in Watts (P)
ENERGY_EFFICIENCY_COEFF = 2e-28# Computation energy efficiency coefficient (kappa)
ENERGY_HARVEST_SCALE_S = 1.0   # Energy harvesting scaling factor (S) [NEW]

# --- Hardware Constraints ---
CPU_FREQ_MAX_HZ = 5e7  # Maximum CPU frequency in Hz (f_max)
CPU_FREQ_MIN_HZ = 1e7  # Minimum CPU frequency in Hz (f_min)
TX_POWER_MAX_W = 50e-6   # Maximum transmit power of WDs in Watts (P_max)
TX_POWER_MIN_W = 0.1e-6  # Minimum transmit power of WDs in Watts (P_min)

# --------------------
# Optimization Weights
# --------------------
TIME_WEIGHT = 1.0     # Weight for total time in the objective function (w_t)
ENERGY_WEIGHT = 1e5   # Weight for total energy in the objective function (w_E)