#!/usr/bin/env python3

import math
import collections

from fractions import Fraction
from typing import Dict, List

import satisfactory


def basic_rate(recipe: Dict) -> Fraction:
    """
    Calculate the rate at which the item is crafted with the default recipe.
    :param recipe:
    :return:
    """
    for output, count in recipe.output.items():
        return Fraction(
            count, recipe.crafting_time
        )


def required_rates(
        recipes: satisfactory.Cookbook,
        remain: List[Dict[str, Fraction]]
) -> Dict[str, Fraction]:
    """
    Calculate the total production rates needed to produce items at the rates
    listed in 'remain' using the provided recipes.
    :param recipes: The source Cookbook
    :param remain: A mapping from requested items to their desired rate.
    :return: A mapping from all required items to their required output rates.
    """

    # Iteratively accumulate a mapping from items to required rates
    required_items = collections.defaultdict(Fraction)

    while remain:
        for dst_name, dst_rate in remain.pop().items():
            # Append the requested item to the output rate. We'll add the
            # inputs just below.
            required_items[dst_name] += dst_rate

            # We can't craft resources, so just ignore them.
            if not recipes.is_component(dst_name):
                assert recipes.is_resource(dst_name)
                continue

            # Assume we're using the default recipe
            dst_recipe = recipes[dst_name].recipes[0]

            # Calculate the fraction of the base rate we need to satisfy, and
            # append our (scaled) inputs to the request queue.
            normal_rate = basic_rate(dst_recipe)
            scale = dst_rate / normal_rate

            for src_name, src_count in dst_recipe.input.items():
                src_rate = Fraction(
                    src_count,
                    dst_recipe.crafting_time
                ) * scale
                remain.append({src_name: src_rate})

    return required_items


# The number of items per minute that each tier of conveyor can carry
conveyor_rates = [0, 60, 120, 270, 480, 780, 900]


def required_machines(recipes: satisfactory.Cookbook, rates: Dict[str, Fraction]) -> Dict[str, int]:
    """
    Calculate the number of machines required to build each item at the
    requested rates. This will, by necessity, round up rates where they are
    below the output rate of the relevant machine.

    :param recipes: The cookbook object we're working from
    :param rates: A mapping from item names to requested output rates.
    :return: A mapping from machine names to counts
    """
    def numberdict():
        return collections.defaultdict(int)
    required_machines = collections.defaultdict(numberdict)

    for name, requested_rate in rates.items():
        if recipes.is_resource(name):
            continue

        descriptor = recipes[name]

        normal_rate = Fraction(
            descriptor.recipes[0].output[name],
            descriptor.recipes[0].crafting_time
        )

        machine = descriptor.machine
        required_machines[machine][name] += requested_rate / normal_rate

    return required_machines


def required_power(recipes: satisfactory.Cookbook, machines: Dict[str, int]) -> int:
    """
    Calculate the cumulative power requirements for a mapping of machine
    names to counts

    :param recipes: The cookbook we're working from
    :param machines: A mapping from machine names to counts
    :return: The total required power in MW
    """
    # Calculate the power requirements for all the machines
    total = 0

    for machine, buckets in machines.items():
        for result, rate in buckets.items():
            count = int(math.ceil(rate))
            source = recipes[machine]
            total += count * source.power_usage
            print(machine, result, math.ceil(rate))

    return total


def plan(recipes: satisfactory.Cookbook, required: Dict[str, Fraction]):
    """
    Print the items and rates, machines and counts, and power requirements
    need to produce the items named in the dict at the mapped rate.

    :param recipes:
    :param required:
    :return:
    """
    rates = required_rates(recipes, required)

    # Note if any particular item is (in aggregate) going to exceed the
    # highest conveyor belt capacity.
    for name, rate in rates.items():
        print(name, rate, float(rate * 60))
        if rate * 60 > conveyor_rates[-1]:
            print("Rate exceeds max conveyor")

    machines = required_machines(recipes, rates)
    power = required_power(recipes, machines)

    print(power, "MW")
    print(math.ceil(power / 150), "fuel generators")


def main():
    recipes = satisfactory.Cookbook('data/recipes')

    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "item",
        nargs="*",
        type=str,
        default=recipes.components(),
        help="The name of an item to produce at full rate"
    )

    args = parser.parse_args()

    # Create an initial name:rate request for all of the target items, then
    # create a plan for their creation.
    request = [{n: basic_rate(recipes[n].recipes[0]) for n in args.item}]

    plan(recipes, request)


if __name__ == '__main__':
    main()