"""
Mass Shell - K11 Mass Coupling for CST Tokens
==============================================

HEX_ID: 43208FA2
AREA: Physical (43)
COMPONENT: Mass (20)
LEVEL: Shell (8F)
INSTANCE: Tertiary (A2)

Core Principle: mass = K11 * (1 + k11_level / 11)
Mass emerges from the K11 hendecagonal constant coupled with position in the manifold.
EM coupling = K5 * K11 * phi_norm gives the electromagnetic mass equivalent.

Celestial Mass Primes:
    Sun Prime: 2179 (central mass reference)
    Earth Prime: 6584 (base reference)
    Jupiter Prime: 2089 (outer mass)
    Moon Prime: 79 (bound mass)

Commands:
    status          - Mass system status
    get HEX         - Get mass of token
    coupling HEX1 HEX2 - Mass coupling between tokens
    momentum HEX    - Momentum vector for token
    center-of-mass  - Compute barycenter

Inputs/Outputs:
    status: None -> Dict (K11, K5, mass_range, celestial_bodies)
    get: hex_id:str -> Dict (mass, k11_level, em_coupling, energy_equivalent)
    coupling: hex1:str, hex2:str -> Dict (gravitational_coupling, reduced_mass, binding_energy)
    momentum: hex_id:str, velocity:float -> Dict (momentum_vector, angular_momentum)
    center_of_mass: hex_ids:list -> Dict (total_mass, center_of_mass coords)
"""

from __future__ import annotations

import math
import json
import argparse
from typing import Dict, Any, Tuple, List, Optional
from datetime import datetime

from .._core.constants import PHI, K5, K11, LIGHT_UNIT, MATRIX_SIZE
from .._core import morton


# =============================================================================
# MASS CONSTANTS
# =============================================================================

# Mass Constants (CST units)
MASS_UNIT = K11              # Base mass quantum
PLANCK_MASS = 2.176434e-8    # kg (for reference)

# Celestial Mass Primes
SUN_PRIME = 2179
EARTH_PRIME = 6584
JUPITER_PRIME = 2089
MOON_PRIME = 79

# Mass Ratios (relative to Earth)
SUN_MASS_RATIO = SUN_PRIME / EARTH_PRIME        # ~0.331
JUPITER_MASS_RATIO = JUPITER_PRIME / EARTH_PRIME # ~0.317
MOON_MASS_RATIO = MOON_PRIME / EARTH_PRIME      # ~0.012

# Energy-Mass equivalence factor
C_SQUARED = PHI * K11 * 131  # CST "c^2" equivalent


# =============================================================================
# CORE FUNCTIONS
# =============================================================================

def hex_to_morton(hex_id: str) -> int:
    """Extract Morton code from hex ID - O(1)."""
    try:
        value = int(hex_id, 16)
        return value & 0xFFFFF
    except ValueError:
        return 0


def morton_decode_wrapped(m: int) -> Tuple[int, int, int]:
    """Decode Morton to wrapped coordinates."""
    x, y, z = morton.decode(m)
    return (x % MATRIX_SIZE, y % MATRIX_SIZE, z % MATRIX_SIZE)


def compute_k11_level(m: int) -> int:
    """Compute K11 level from Morton code - O(1)."""
    return (m % 23) - 11  # Range: [-11, +11]


def compute_phi_norm(coords: Tuple[int, int, int]) -> float:
    """Compute phi-normalized distance from center - O(1)."""
    center = MATRIX_SIZE // 2
    dx = coords[0] - center
    dy = coords[1] - center
    dz = coords[2] - center
    distance = math.sqrt(dx*dx + dy*dy + dz*dz)
    return distance / (MATRIX_SIZE * PHI)


def compute_mass(hex_id: str) -> float:
    """
    Compute mass of a token from K11 coupling.

    mass = K11 * (1 + k11_level / 11)

    Range: K11 * 0 to K11 * 2 = [0, 18.408]
    """
    m = hex_to_morton(hex_id)
    k11_level = compute_k11_level(m)
    mass = K11 * (1 + k11_level / 11)
    return max(0, mass)  # Mass cannot be negative


def compute_em_coupling(hex_id: str) -> float:
    """
    Compute electromagnetic coupling.

    em_coupling = K5 * K11 * phi_norm
    """
    m = hex_to_morton(hex_id)
    coords = morton_decode_wrapped(m)
    phi_norm = compute_phi_norm(coords)
    return K5 * K11 * phi_norm


def compute_energy_equivalent(mass: float) -> float:
    """
    E = m * c^2 (CST equivalent)
    """
    return mass * C_SQUARED


# =============================================================================
# MASS OPERATIONS
# =============================================================================

def get_token_mass(hex_id: str) -> Dict[str, Any]:
    """Get complete mass data for a token - O(1)."""
    m = hex_to_morton(hex_id)
    coords = morton_decode_wrapped(m)
    k11_level = compute_k11_level(m)
    phi_norm = compute_phi_norm(coords)

    mass = compute_mass(hex_id)
    em_coupling = compute_em_coupling(hex_id)
    energy = compute_energy_equivalent(mass)

    # Basin as "isotope" indicator
    basin = m % LIGHT_UNIT

    return {
        "hex_id": hex_id,
        "morton": m,
        "coords": coords,
        "k11_level": k11_level,
        "phi_norm": round(phi_norm, 6),
        "mass": round(mass, 6),
        "em_coupling": round(em_coupling, 6),
        "energy_equivalent": round(energy, 4),
        "basin": basin,
        "mass_classification": classify_mass(mass),
        "computed_at": datetime.now().isoformat()
    }


def classify_mass(mass: float) -> str:
    """Classify mass into categories."""
    if mass < K11 * 0.5:
        return "LIGHT"       # Below half K11
    elif mass < K11:
        return "STANDARD"    # Around K11
    elif mass < K11 * 1.5:
        return "HEAVY"       # Above K11
    else:
        return "MASSIVE"     # Near maximum


def compute_mass_coupling(hex1: str, hex2: str) -> Dict[str, Any]:
    """
    Compute gravitational mass coupling between two tokens.

    coupling = (m1 * m2) / d^2

    Uses CST distance and mass values.
    """
    m1_data = get_token_mass(hex1)
    m2_data = get_token_mass(hex2)

    c1 = m1_data["coords"]
    c2 = m2_data["coords"]

    # Distance
    dx = c1[0] - c2[0]
    dy = c1[1] - c2[1]
    dz = c1[2] - c2[2]
    distance = math.sqrt(dx*dx + dy*dy + dz*dz)

    # Avoid division by zero
    if distance < 0.1:
        distance = 0.1

    # Gravitational coupling
    mass1 = m1_data["mass"]
    mass2 = m2_data["mass"]
    coupling = (mass1 * mass2) / (distance * distance)

    # EM coupling sum
    em_total = m1_data["em_coupling"] + m2_data["em_coupling"]

    # Reduced mass
    if mass1 + mass2 > 0:
        reduced_mass = (mass1 * mass2) / (mass1 + mass2)
    else:
        reduced_mass = 0

    return {
        "hex1": hex1,
        "hex2": hex2,
        "mass1": round(mass1, 6),
        "mass2": round(mass2, 6),
        "distance": round(distance, 4),
        "gravitational_coupling": round(coupling, 6),
        "em_coupling_sum": round(em_total, 6),
        "reduced_mass": round(reduced_mass, 6),
        "binding_energy": round(coupling * C_SQUARED * 0.01, 4),
        "computed_at": datetime.now().isoformat()
    }


def compute_momentum(hex_id: str, velocity: float = 1.0) -> Dict[str, Any]:
    """
    Compute momentum vector for a token.

    p = m * v

    Velocity direction derived from orbital position.
    """
    mass_data = get_token_mass(hex_id)
    coords = mass_data["coords"]

    # Velocity direction from center (orbital tangent)
    center = MATRIX_SIZE // 2
    dx = coords[0] - center
    dy = coords[1] - center
    dz = coords[2] - center
    r = math.sqrt(dx*dx + dy*dy + dz*dz)

    if r > 0:
        # Tangent vector (perpendicular to radial)
        vx = -dy / r * velocity
        vy = dx / r * velocity
        vz = 0.0
    else:
        vx = vy = vz = 0.0

    mass = mass_data["mass"]
    px = mass * vx
    py = mass * vy
    pz = mass * vz
    p_magnitude = math.sqrt(px*px + py*py + pz*pz)

    # Angular momentum L = r x p
    lx = dy * pz - dz * py
    ly = dz * px - dx * pz
    lz = dx * py - dy * px
    l_magnitude = math.sqrt(lx*lx + ly*ly + lz*lz)

    return {
        "hex_id": hex_id,
        "mass": round(mass, 6),
        "velocity": velocity,
        "velocity_vector": (round(vx, 4), round(vy, 4), round(vz, 4)),
        "momentum_vector": (round(px, 4), round(py, 4), round(pz, 4)),
        "momentum_magnitude": round(p_magnitude, 6),
        "angular_momentum_vector": (round(lx, 4), round(ly, 4), round(lz, 4)),
        "angular_momentum_magnitude": round(l_magnitude, 6),
        "kinetic_energy": round(0.5 * mass * velocity * velocity, 6),
        "computed_at": datetime.now().isoformat()
    }


def compute_center_of_mass(hex_ids: List[str] = None) -> Dict[str, Any]:
    """
    Compute center of mass for a set of tokens.
    If no tokens provided, uses celestial primes.
    """
    if hex_ids is None:
        # Default: celestial bodies
        hex_ids = [
            f"{SUN_PRIME:08X}",
            f"{EARTH_PRIME:08X}",
            f"{MOON_PRIME:08X}",
            f"{JUPITER_PRIME:08X}"
        ]

    total_mass = 0.0
    weighted_x = 0.0
    weighted_y = 0.0
    weighted_z = 0.0
    masses = []

    for hex_id in hex_ids:
        data = get_token_mass(hex_id)
        m = data["mass"]
        c = data["coords"]
        masses.append({"hex_id": hex_id, "mass": m, "coords": c})
        total_mass += m
        weighted_x += m * c[0]
        weighted_y += m * c[1]
        weighted_z += m * c[2]

    if total_mass > 0:
        com_x = weighted_x / total_mass
        com_y = weighted_y / total_mass
        com_z = weighted_z / total_mass
    else:
        com_x = com_y = com_z = MATRIX_SIZE // 2

    return {
        "token_count": len(hex_ids),
        "total_mass": round(total_mass, 6),
        "center_of_mass": (round(com_x, 4), round(com_y, 4), round(com_z, 4)),
        "masses": masses,
        "computed_at": datetime.now().isoformat()
    }


def get_system_status() -> Dict[str, Any]:
    """Get mass shell system status."""

    # Compute celestial body masses
    celestial = {}
    for name, prime in [("sun", SUN_PRIME), ("earth", EARTH_PRIME),
                        ("jupiter", JUPITER_PRIME), ("moon", MOON_PRIME)]:
        hex_id = f"{prime:08X}"
        data = get_token_mass(hex_id)
        celestial[name] = {
            "prime": prime,
            "hex_id": hex_id,
            "mass": data["mass"],
            "k11_level": data["k11_level"],
            "basin": data["basin"]
        }

    return {
        "shell": "mass_shell",
        "version": "1.0.0",
        "hex_id": "43208FA2",
        "status": "ACTIVE",
        "constants": {
            "K11": K11,
            "K5": K5,
            "PHI": PHI,
            "MASS_UNIT": MASS_UNIT,
            "C_SQUARED": round(C_SQUARED, 4)
        },
        "mass_formula": "mass = K11 * (1 + k11_level / 11)",
        "em_formula": "em_coupling = K5 * K11 * phi_norm",
        "mass_range": f"[0, {round(K11 * 2, 4)}]",
        "celestial_bodies": celestial,
        "timestamp": datetime.now().isoformat()
    }


# =============================================================================
# SHELL CLASS
# =============================================================================

class MassShell:
    """Mass Shell - K11 Mass Coupling for CST Tokens."""

    def __init__(self):
        self.name = "mass"
        self.hex_id = "43208FA2"
        self.version = "1.0.0"

    def status(self) -> Dict[str, Any]:
        """Get mass system status."""
        return get_system_status()

    def get(self, hex_id: str) -> Dict[str, Any]:
        """Get mass of token."""
        return get_token_mass(hex_id)

    def coupling(self, hex1: str, hex2: str) -> Dict[str, Any]:
        """Compute mass coupling between tokens."""
        return compute_mass_coupling(hex1, hex2)

    def momentum(self, hex_id: str, velocity: float = 1.0) -> Dict[str, Any]:
        """Compute momentum vector for token."""
        return compute_momentum(hex_id, velocity)

    def center_of_mass(self, hex_ids: List[str] = None) -> Dict[str, Any]:
        """Compute center of mass."""
        return compute_center_of_mass(hex_ids)

    def execute(self, command: str, args: Dict[str, Any] = None) -> Dict[str, Any]:
        """Execute shell command."""
        args = args or {}

        if command == "status":
            return self.status()
        elif command == "get":
            return self.get(args.get("hex_id", "2A9E0100"))
        elif command == "coupling":
            return self.coupling(
                args.get("hex1", "2A9E0100"),
                args.get("hex2", "43208FAC")
            )
        elif command == "momentum":
            return self.momentum(
                args.get("hex_id", "2A9E0100"),
                args.get("velocity", 1.0)
            )
        elif command == "center-of-mass":
            return self.center_of_mass(args.get("tokens"))
        else:
            return {"error": f"Unknown command: {command}"}

    def run(self, command: str, *args, as_json: bool = False) -> Any:
        """Run a shell command (CLI interface)."""
        commands = {
            'status': lambda: self.status(),
            'get': lambda: self.get(args[0] if args else "2A9E0100"),
            'coupling': lambda: self.coupling(args[0] if len(args) > 0 else "2A9E0100",
                                               args[1] if len(args) > 1 else "43208FAC"),
            'momentum': lambda: self.momentum(args[0] if len(args) > 0 else "2A9E0100",
                                               float(args[1]) if len(args) > 1 else 1.0),
            'center-of-mass': lambda: self.center_of_mass(list(args) if args else None),
        }
        if command in commands:
            result = commands[command]()
            if as_json and isinstance(result, dict):
                return json.dumps(result, indent=2, default=str)
            return result
        return {'error': f'Unknown command: {command}', 'available': list(commands.keys())}


# =============================================================================
# CLI INTERFACE
# =============================================================================

def main():
    parser = argparse.ArgumentParser(
        description='Mass Shell - K11 Mass Coupling for CST Tokens',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python -m cubesquare.shells.mass --status
    python -m cubesquare.shells.mass --get 2A9E0100
    python -m cubesquare.shells.mass --coupling 2A9E0100 43208FAC
    python -m cubesquare.shells.mass --momentum 2A9E0100
    python -m cubesquare.shells.mass --center-of-mass

Core Formula: mass = K11 * (1 + k11_level / 11)
"""
    )

    parser.add_argument('--status', action='store_true', help='Mass system status')
    parser.add_argument('--get', metavar='HEX', help='Get mass of token')
    parser.add_argument('--coupling', nargs=2, metavar=('HEX1', 'HEX2'), help='Mass coupling')
    parser.add_argument('--momentum', metavar='HEX', help='Momentum vector')
    parser.add_argument('--velocity', type=float, default=1.0, help='Velocity (default: 1.0)')
    parser.add_argument('--center-of-mass', action='store_true', help='Compute barycenter')
    parser.add_argument('--tokens', nargs='+', metavar='HEX', help='Tokens for center-of-mass')
    parser.add_argument('--json', action='store_true', help='JSON output')

    args = parser.parse_args()

    shell = MassShell()
    result = None

    if args.status:
        result = shell.status()
    elif args.get:
        result = shell.get(args.get)
    elif args.coupling:
        result = shell.coupling(args.coupling[0], args.coupling[1])
    elif args.momentum:
        result = shell.momentum(args.momentum, args.velocity)
    elif args.center_of_mass:
        result = shell.center_of_mass(args.tokens)
    else:
        parser.print_help()
        return

    if result:
        print(json.dumps(result, indent=2, default=str))


if __name__ == '__main__':
    main()
