import random
import sys
from collections import deque

# Graph of physical formulas
formulas = [
    {"name": "W=P·t", "vars": ["W", "P", "t"]},
    {"name": "P=U·I", "vars": ["P", "U", "I"]},
    {"name": "Q=I·t", "vars": ["Q", "I", "t"]},
    {"name": "U=I·R", "vars": ["U", "I", "R"]},
    {"name": "W=Q·U", "vars": ["W", "Q", "U"]},
]

all_vars = ["P", "W", "t", "U", "I", "Q", "R"]

# Translations
translations = {
    "en": {
        "intro": "A heater is connected to a constant voltage source. A constant current flows through it.",
        "given": "Given",
        "find": "Find",
        "solution_scheme": "Solution scheme",
        "step": "Step",
        "substitute": "substitute",
        "into_formula": "into formula",
        "calculate": "calculate",
        "and": "and",
        "var_names": {
            "P": "power",
            "W": "energy",
            "t": "time",
            "U": "voltage",
            "I": "current",
            "Q": "charge",
            "R": "resistance",
        },
    },
    "pl": {
        "intro": "Do źródła stałego napięcia podłączono grzałkę. Płynie przez nią stały prąd.",
        "given": "Dane",
        "find": "Szukane",
        "solution_scheme": "Schemat rozwiązania",
        "step": "Krok",
        "substitute": "wstaw",
        "into_formula": "do wzoru",
        "calculate": "wylicz",
        "and": "oraz",
        "var_names": {
            "P": "moc",
            "W": "energia",
            "t": "czas",
            "U": "napięcie",
            "I": "natężenie prądu",
            "Q": "ładunek",
            "R": "opór",
        },
    },
}

units = {
    "P": "W",
    "W": "J",
    "t": "s",
    "U": "V",
    "I": "A",
    "Q": "C",
    "R": "Ω",
}

def find_path(known_vars, target_var):
    """Find shortest path from known variables to target variable."""
    if target_var in known_vars:
        return []
    
    queue = deque([(set(known_vars), [])])
    visited = set()
    
    while queue:
        current_known, path = queue.popleft()
        
        state = tuple(sorted(current_known))
        if state in visited:
            continue
        visited.add(state)
        
        if target_var in current_known:
            return path
        
        for formula in formulas:
            fvars = formula["vars"]
            known_in_formula = [v for v in fvars if v in current_known]
            unknown_in_formula = [v for v in fvars if v not in current_known]
            
            if len(known_in_formula) >= 2 and len(unknown_in_formula) == 1:
                new_var = unknown_in_formula[0]
                new_known = current_known | {new_var}
                new_path = path + [(formula, known_in_formula, new_var)]
                queue.append((new_known, new_path))
    
    return None

def generate_random_task(n_steps=None):
    """Generate random task with random number of known variables and 1 target variable."""
    min_known = 2
    max_known = len(all_vars) - 1
    n_known = random.randint(min_known, max_known)
    
    known_vars = random.sample(all_vars, n_known)
    remaining = [v for v in all_vars if v not in known_vars]
    target_var = random.choice(remaining)
    
    path = find_path(known_vars, target_var)
    
    if path is None:
        return None
    
    if n_steps is not None and len(path) != n_steps:
        return None
    
    # Filter out unused variables
    used_vars = set()
    for formula, known_in_formula, new_var in path:
        for v in known_in_formula:
            used_vars.add(v)
    
    actual_known = [v for v in known_vars if v in used_vars]
    
    return actual_known, target_var, path

def generate_task_text(lang="en", n_steps=None):
    """Generate complete task with numerical values and solution."""
    max_attempts = 1000 if n_steps is not None else 100
    
    for _ in range(max_attempts):
        result = generate_random_task(n_steps)
        if result is not None:
            break
    else:
        if n_steps is not None:
            return f"Cannot generate task with exactly {n_steps} steps after {max_attempts} attempts."
        return "Cannot generate task for these parameters."
    
    known_vars, target_var, path = result
    t = translations[lang]
    
    # Generate random values for known variables
    values = {}
    for var in known_vars:
        if var == "t":
            values[var] = random.randint(5, 60)
        elif var == "R":
            values[var] = random.randint(10, 100)
        elif var == "I":
            values[var] = random.randint(1, 10)
        elif var == "U":
            values[var] = random.randint(10, 240)
        elif var == "Q":
            values[var] = random.randint(50, 500)
        elif var == "P":
            values[var] = random.randint(100, 2000)
        elif var == "W":
            values[var] = random.randint(1000, 10000)
    
    # Calculate target variable value
    for formula, known_in_formula, new_var in path:
        fname = formula["name"]
        
        if fname == "W=P·t":
            if new_var == "W":
                values["W"] = values["P"] * values["t"]
            elif new_var == "P":
                values["P"] = values["W"] / values["t"]
            else:
                values["t"] = values["W"] / values["P"]
        elif fname == "P=U·I":
            if new_var == "P":
                values["P"] = values["U"] * values["I"]
            elif new_var == "U":
                values["U"] = values["P"] / values["I"]
            else:
                values["I"] = values["P"] / values["U"]
        elif fname == "Q=I·t":
            if new_var == "Q":
                values["Q"] = values["I"] * values["t"]
            elif new_var == "I":
                values["I"] = values["Q"] / values["t"]
            else:
                values["t"] = values["Q"] / values["I"]
        elif fname == "U=I·R":
            if new_var == "U":
                values["U"] = values["I"] * values["R"]
            elif new_var == "I":
                values["I"] = values["U"] / values["R"]
            else:
                values["R"] = values["U"] / values["I"]
        elif fname == "W=Q·U":
            if new_var == "W":
                values["W"] = values["Q"] * values["U"]
            elif new_var == "Q":
                values["Q"] = values["W"] / values["U"]
            else:
                values["U"] = values["W"] / values["Q"]
    
    # Format task
    output = f"{t['intro']}\n\n"
    output += f"{t['given']}:\n"
    for var in known_vars:
        output += f"  {t['var_names'][var]}: {values[var]} {units[var]}\n"
    output += f"\n{t['find']}: {t['var_names'][target_var]} ({values[target_var]:.2f} {units[target_var]})\n"
    
    output += f"\n--- {t['solution_scheme']} ---\n"
    for i, (formula, known_in_formula, new_var) in enumerate(path, 1):
        known_str = f" {t['and']} ".join(known_in_formula)
        output += f"{t['step']} {i}: {t['into_formula']} \"{formula['name']}\" {t['substitute']} {known_str}, {t['calculate']} {new_var}.\n"
    
    return output

def main():
    # Parse command line arguments
    num_tasks = 1
    lang = "en"
    num_steps = None
    
    args = sys.argv[1:]
    i = 0
    while i < len(args):
        if args[i] in ["-n", "--num-tasks"]:
            num_tasks = int(args[i + 1])
            i += 2
        elif args[i] in ["-l", "--lang"]:
            lang = args[i + 1]
            if lang not in ["en", "pl"]:
                print(f"Error: Unsupported language '{lang}'. Use 'en' or 'pl'.")
                sys.exit(1)
            i += 2
        elif args[i] in ["-s", "--steps"]:
            num_steps = int(args[i + 1])
            i += 2
        elif args[i] in ["-h", "--help"]:
            print("Usage: python3 fizyka.py [options]")
            print("\nOptions:")
            print("  -n, --num-tasks N    Number of tasks to generate (default: 1)")
            print("  -l, --lang LANG      Language: 'en' or 'pl' (default: en)")
            print("  -s, --steps N        Number of solution steps (default: any)")
            print("  -h, --help           Show this help message")
            sys.exit(0)
        else:
            print(f"Error: Unknown argument '{args[i]}'")
            print("Use -h or --help for usage information.")
            sys.exit(1)
    
    # Generate tasks
    for i in range(num_tasks):
        print(generate_task_text(lang, num_steps))
        if i < num_tasks - 1:
            print("\n" + "="*60 + "\n")

if __name__ == "__main__":
    main()

