import matplotlib.pyplot as plt
import graphviz

# Define the structure of the diagram

def standard_production_function():

    dot = graphviz.Digraph()

    dot.attr(rankdir='LR')

    # Add nodes and edges
    dot.node('Output', 'Output (CES Aggregator)', shape='Mrecord', color='blue', fillcolor='lightblue', style='filled')

    # Energy Composite
    dot.node('Energy Composite', 'Energy Composite (CES Aggregator)', shape='Mrecord', color='red', fillcolor='pink', style='filled')
    dot.edge('Output', 'Energy Composite', color='blue', dir='back')

    dot.node('Electricity', 'Electricity (CES Aggregator)', shape='Mrecord', color='orange', fillcolor='orange', style='filled')
    dot.node('Coal', 'Coal', color='red')
    dot.node('Gas', 'Gas', color='red')
    dot.node('Oil', 'Oil', color='red')
    dot.node('Petroleum', 'Petroleum', color='red')
    dot.edge('Energy Composite', 'Electricity', color='red', dir='back')
    dot.edge('Energy Composite', 'Coal', color='red', dir='back')
    dot.edge('Energy Composite', 'Gas', color='red', dir='back')
    dot.edge('Energy Composite', 'Oil', color='red', dir='back')
    dot.edge('Energy Composite', 'Petroleum', color='red', dir='back')

    # Capital
    dot.node('Capital', 'Capital', color='blue')
    dot.edge('Output', 'Capital', color='blue', dir='back')

    # Labor
    dot.node('Labor', 'Labor', color='blue')
    dot.edge('Output', 'Labor', color='blue', dir='back')

    # Material Composite
    dot.node('Material Composite', 'Material Composite (CES Aggregator)', shape='Mrecord', color='green', fillcolor='lightgreen', style='filled')
    dot.edge('Output', 'Material Composite', color='blue', dir='back')

    dot.node('Construction', 'Construction', color='green')
    dot.node('Mining', 'Other mining', color='green')
    dot.node('Agriculture', 'Agriculture', color='green')
    dot.node('Durable manufacturing', 'Durable manufacturing', color='green')
    dot.node('Non-durable manufacturing', 'Non-durable manufacturing', color='green')
    dot.node('Services', 'Services', color='green')

    dot.edge('Material Composite', 'Construction', color='green', dir='back')
    dot.edge('Material Composite', 'Mining', color='green', dir='back')
    dot.edge('Material Composite', 'Agriculture', color='green', dir='back')
    dot.edge('Material Composite', 'Durable manufacturing', color='green', dir='back')
    dot.edge('Material Composite', 'Non-durable manufacturing', color='green', dir='back')
    dot.edge('Material Composite', 'Services', color='green', dir='back')

    # Generate the diagram as an SVG
    svg_path = 'standard_production_function'
    dot.render(svg_path, format='svg', cleanup=True)

    svg_path


def electricity_distribution_production_function():

    dot = graphviz.Digraph()

    dot.attr(rankdir='LR')

    # Add nodes and edges
    dot.node('Electricity', 'Electricity distribution (CES Aggregator)', shape='Mrecord', color='orange', fillcolor='orange', style='filled')

    dot.node('Coal generation', 'Coal generation', color='orange')
    dot.node('Oil generation', 'Oil generation', color='orange')
    dot.node('Gas generation', 'Gas generation', color='orange')
    dot.node('Nuclear generation', 'Nuclear generation', color='orange')
    dot.node('Solar generation', 'Solar generation', color='orange')
    dot.node('Wind generation', 'Wind generation', color='orange')
    dot.node('Hydro generation', 'Hydro generation', color='orange')
    dot.node('Other generation', 'Other generation', color='orange')

    dot.edge('Electricity', 'Coal generation', color='orange', dir='back')
    dot.edge('Electricity', 'Oil generation', color='orange', dir='back')
    dot.edge('Electricity', 'Gas generation', color='orange', dir='back')
    dot.edge('Electricity', 'Nuclear generation', color='orange', dir='back')
    dot.edge('Electricity', 'Solar generation', color='orange', dir='back')
    dot.edge('Electricity', 'Wind generation', color='orange', dir='back')
    dot.edge('Electricity', 'Hydro generation', color='orange', dir='back')
    dot.edge('Electricity', 'Other generation', color='orange', dir='back')

    # Generate the diagram as an SVG
    svg_path = 'electricity_distribution_production_function'
    dot.render(svg_path, format='svg', cleanup=True)

    svg_path

standard_production_function()

electricity_distribution_production_function()