# sankey_dha.2.py
# Requirements: plotly
# Install with: pip install plotly

import plotly.graph_objects as go

labels = [
    "Microalgae (A=200)",
    "Algal farming (AF=30)",
    "Zooplankton",
    "Krill",
    "Small forage fish",
    "Large predatory fish",
    "Fishmeal/Fishoil (FM=25)",
    "Aquaculture feed (aqu_feed=20)",
    "Farmed seafood (H_farmed=18)",
    "Wild-caught seafood (H_wild=50)",
    "Supplements (H_supp=10)",
    "Human total intake (H_total=78)",
    "Losses / detritus / benthos"
]

idx = {label: i for i, label in enumerate(labels)}

# Flows derived from the mass-balance above
sources = [
    idx["Microalgae (A=200)"],  # microalgae -> zooplankton
    idx["Microalgae (A=200)"],  # microalgae -> krill
    idx["Microalgae (A=200)"],  # microalgae -> losses
    idx["Zooplankton"],         # zooplankton -> krill
    idx["Zooplankton"],         # zooplankton -> small fish
    idx["Zooplankton"],         # zooplankton -> losses
    idx["Krill"],               # krill -> small fish
    idx["Krill"],               # krill -> losses
    idx["Small forage fish"],   # small fish -> large fish
    idx["Small forage fish"],   # small fish -> fishmeal
    idx["Small forage fish"],   # small fish -> direct human (small)
    idx["Small forage fish"],   # small fish -> losses
    idx["Large predatory fish"],# large fish -> wild-caught human
    idx["Large predatory fish"],# large fish -> losses
    idx["Fishmeal/Fishoil (FM=25)"], # FM -> aquaculture feed (15)
    idx["Algal farming (AF=30)"],    # AF -> aquaculture feed (5)
    idx["Algal farming (AF=30)"],    # AF -> supplements (10)
    idx["Fishmeal/Fishoil (FM=25)"], # FM remainder -> storage/other (10)
    idx["Aquaculture feed (aqu_feed=20)"], # aquaculture feed -> farmed seafood
    idx["Aquaculture feed (aqu_feed=20)"], # aquaculture feed -> aquaculture losses
    idx["Farmed seafood (H_farmed=18)"],   # farmed seafood -> human total
    idx["Wild-caught seafood (H_wild=50)"],# wild-caught -> human total
    idx["Supplements (H_supp=10)"],        # supplements -> human total
    # route losses to losses node for clarity
]

targets = [
    idx["Zooplankton"],       # microalgae -> zooplankton
    idx["Krill"],             # microalgae -> krill
    idx["Losses / detritus / benthos"], # microalgae -> losses
    idx["Krill"],             # zooplankton -> krill
    idx["Small forage fish"], # zooplankton -> small fish
    idx["Losses / detritus / benthos"], # zooplankton -> losses
    idx["Small forage fish"], # krill -> small fish
    idx["Losses / detritus / benthos"], # krill -> losses
    idx["Large predatory fish"], # small fish -> large fish
    idx["Fishmeal/Fishoil (FM=25)"], # small fish -> fishmeal
    idx["Wild-caught seafood (H_wild=50)"], # small fish -> direct human
    idx["Losses / detritus / benthos"],    # small fish -> losses
    idx["Wild-caught seafood (H_wild=50)"], # large fish -> wild-caught humans
    idx["Losses / detritus / benthos"],    # large fish -> losses
    idx["Aquaculture feed (aqu_feed=20)"], # FM -> aquaculture feed (15)
    idx["Aquaculture feed (aqu_feed=20)"], # AF -> aquaculture feed (5)
    idx["Supplements (H_supp=10)"],        # AF -> supplements (10)
    idx["Losses / detritus / benthos"],    # FM remainder
    idx["Farmed seafood (H_farmed=18)"],   # aquaculture feed -> farmed seafood
    idx["Losses / detritus / benthos"],    # aquaculture feed -> losses
    idx["Human total intake (H_total=78)"],# farmed seafood -> human total
    idx["Human total intake (H_total=78)"],# wild-caught -> human total
    idx["Human total intake (H_total=78)"],# supplements -> human total
]

values = [
    90,  # microalgae -> zooplankton (a1)
    70,  # microalgae -> krill (a2)
    40,  # microalgae -> losses (L_a)
    30,  # zooplankton -> krill (z1)
    40,  # zooplankton -> small fish (z2)
    20,  # zooplankton -> losses (L_z)
    60,  # krill -> small fish (k1)
    40,  # krill -> losses (L_k)
    60,  # small fish -> large fish (s1)
    25,  # small fish -> fishmeal/fishoil (s2)
    5,   # small fish -> direct human (s3)
    10,  # small fish -> losses (L_s)
    50,  # large fish -> wild-caught humans (p1)
    10,  # large fish -> losses (L_p)
    15,  # FM -> aquaculture feed (part of 25)
    5,   # AF -> aquaculture feed (part of 30)
    10,  # AF -> supplements
    10,  # FM remainder -> storage/other
    18,  # aquaculture feed -> farmed seafood (H_farmed)
    2,   # aquaculture feed -> losses
    18,  # farmed seafood -> human total
    50,  # wild-caught -> human total
    10   # supplements -> human total
]

link_colors = ["rgba(44,160,44,0.6)"] * len(values)

fig = go.Figure(go.Sankey(
    arrangement="snap",
    node=dict(
        pad=15,
        thickness=18,
        line=dict(color="black", width=0.5),
        label=labels,
        color=["rgba(20,120,180,0.8)",
               "rgba(70,130,180,0.8)",
               "rgba(100,160,100,0.8)",
               "rgba(100,200,150,0.8)",
               "rgba(200,160,100,0.8)",
               "rgba(200,100,100,0.8)",
               "rgba(150,120,120,0.8)",
               "rgba(180,140,80,0.8)",
               "rgba(120,100,200,0.8)",
               "rgba(80,160,200,0.8)",
               "rgba(120,200,180,0.8)",
               "rgba(200,120,160,0.8)",
               "rgba(160,160,160,0.8)"]
    ),
    link=dict(
        source=sources,
        target=targets,
        value=values,
        color=link_colors
    )
))

fig.update_layout(
    title_text="Sankey: Consistent DHA Sources and Sinks (illustrative steady-state flows)",
    font_size=12,
    width=1100,
    height=650
)

# write standalone HTML file
fig.write_html("sankey_dha.2.html", full_html=True, include_plotlyjs="cdn")

# also provide HTML string if you want to embed elsewhere
html_str = fig.to_html(full_html=True, include_plotlyjs="cdn")
# optionally write html_str to a file as shown above