# llm_training_estimator.py

def get_gpu_teraflops(gpu_type: str, exo_flops: float = None) -> float:
    """
    Returns TFLOPs/s based on GPU type or Exo input.
    """
    gpu_lookup = {
        "A100": 312,   # FP16
        "H100": 700,
    }

    if gpu_type == "Exo":
        if exo_flops is None or exo_flops <= 0:
            raise ValueError("Exo TFLOPs must be provided.")
        return exo_flops

    return gpu_lookup.get(gpu_type, 0)


def get_gpu_cost_per_tflop_hour(gpu_type: str, manual_cost: float = None) -> float:
    """
    Returns ₹ cost per TFLOP-hour based on GPU type.
    """
    cost_lookup = {
        "A100": 17.64 * 83,   # $17.64 × ₹83 (approx.) = ₹1464 per TFLOP/hr
        "H100": 7.56 * 83     # ₹627 per TFLOP/hr
    }

    if gpu_type == "Exo" and manual_cost:
        return manual_cost
    return cost_lookup.get(gpu_type, 0)


def estimate_training_time_and_cost(params_billion: float, tokens_billion: float, teraflops: float,
                                    cost_per_tflop_hr: float = 0.0):
    """
    Estimate training time and cost.

    Returns:
    - FLOPs required
    - Time (hours/days)
    - Cost (₹)
    """
    total_flops_required = 6 * params_billion * 1e9 * tokens_billion * 1e9  # 6 × N × T
    time_seconds = total_flops_required / (teraflops * 1e12)
    time_hours = time_seconds / 3600
    time_days = time_hours / 24

    total_cost = (teraflops * time_hours) * cost_per_tflop_hr if cost_per_tflop_hr > 0 else 0

    return {
        "flops_required": total_flops_required,
        "time_hours": time_hours,
        "time_days": time_days,
        "total_cost": total_cost
    }