from sympy import *
from sympy.utilities.codegen import RustCodeGen


def to_latex(eq, inline=False):
    mode = "equation"
    itex = True
    if inline:
        mode = "inline"
        itex = False
    return latex(eq, mul_symbol=" \\cdot ", mode=mode, itex=itex)


def solve_for(eq, sym):
    return Eq(sym, solve(eq, sym)[0])


init_printing(use_latex=True, latex_mode="equation", mul_symbol="\\cdot")

var("m_ship m_fuel m_opt f_useable f_max l p boost dist dist_max Fuel B_g e_fuel")

mass = m_ship + m_fuel  # total mass of ship+fuel
m_opt = m_opt * boost  # supercharging increases optimized mass

available_fuel = Min(
    f_max, m_fuel
)  # limit maximum fuel consumption to FSD max fuel limit

eq_fuel = Eq(Fuel, l * 0.001 * (((dist * mass) / m_opt) ** p))  # FSD Fuel equation



eq_fuel_boost = eq_fuel.subs({"dist":dist+B_g,"Fuel":available_fuel*e_fuel}) # FSD Booster boosts maximum distance by B_g
eq_d_boost = solve_for(eq_fuel_boost, dist)  # solve for distance
# eq_d_boost = eq_d_boost.subs({"Fuel":f_max,"m_fuel":f_max})  # Assume maximum jump range

print(to_latex(eq_d_boost))
print(to_latex(solve_for(eq_d_boost,e_fuel)))

exit()

max_range = eq_d.subs(
    {m_fuel: f_max, dist: dist_max}
)  # Compute maximum jump range by assuming f_max tons of fuel in tank

full_eq = eq_d.subs(
    Min(f_max, m_fuel), Min(f_max, m_fuel) * fuel_mult
).subs(
    max_range.lhs, max_range.rhs
)  # substitute everything in

docs = [
    (
        eq_fuel,
        "FSD Fuel consumption ([E:D Wiki](https://elite-dangerous.fandom.com/wiki/Frame_Shift_Drive#Hyperspace_Fuel_Equation)):",
    ),
    (
        eq_d,
        "Solving for $dist$ gives the jump range (in Ly) for a given amount of fuel (in tons) as:",
    ),
    (
        max_range,
        "Assuming $f_{max}$ tons of available fuel gives us the maximum jump range for a single jump as:",
    ),
    (
        fuel_mult,
        "Since the guardian FSD booster increases the maximum jump range by $B_g$ Ly we can calculate a correction factor for the fuel consumption as:",
    ),
    (
        eq_d.subs(Min(f_max, m_fuel), Min(f_max, m_fuel) * e_fuel),
        "Incorporating $e_{fuel}$ into the distance equation yields",
    ),
    (
        eq_d.subs(Min(f_max, m_fuel), Min(f_max, m_fuel) * fuel_mult),
        "Expanding $e_{fuel}$ yields",
    ),
    (full_eq, "Finally, Expanding $dist_{max}$ yields the full equation as"),
]


var_defs = [
    ("Fuel", "is the fuel needed to jump (in tons)"),
    ("l", "is the linear constant of your FSD (depends on the rating)"),
    ("p", "is the power constant of your FSD (depends on the class)"),
    ("m_ship", "is the mass of your ship (including cargo)"),
    ("m_fuel", "is the amount of fuel in tons currently stored in your tanks"),
    ("m_opt", "is the optimized mass of your FSD (in tons)"),
    ("f_max", "is the maximum amount of fuel your FSD can use per jump"),
    (
        "boost",
        'is the "boost factor" of your FSD (1.0 when jumping normally, 1.5 when supercharged by a white dwarf, 4.0 for a neutron star, etc)',
    ),
    ("dist", "is the distance you can jump with a given fuel amount"),
    ("dist_max", "is the maximum distance you can jump (when $m_{fuel}=f_{max}$)"),
    ("B_g", "is the amount of Ly added by your Guardian FSD Booster"),
    ("e_fuel", "is the efficiency increase added by the Guardian FSD Booster"),
]
for eq, doc in docs:
    eq=simplify(eq)
    if doc:
        print(doc, to_latex(eq))
    else:
        print(to_latex(eq))
    print()

print("Where:")
for name, desc in var_defs:
    print("- {} {}".format(to_latex(symbols(name), True), desc))