import sciline
from typing import NewType

Bank = NewType("Bank", str)
Counts = NewType("Counts", int)
Result = NewType("Result", float)
Param = NewType("Param", float)
Normalized = NewType("Normalized", float)


def load(bank: Bank) -> Counts:
    return Counts(123 if bank == "bank1" else 456)


def process(counts: Counts, param: Param) -> Result:
    return Result(counts * param)


def normalize(result: Result) -> Normalized:
    return Normalized(result / 10)


pl = sciline.Pipeline((load, process, normalize))
pl[Param] = Param(0.5)
pl[Bank] = Bank("bank1")
graph = pl.get(Normalized)
graph.visualize(graph_attr={"rankdir": "LR"})

import cyclebane as cb
import pandas as pd

df = pd.DataFrame({Bank: ["bank1", "bank2"]}).rename_axis(index="bank")
params = pd.DataFrame({Param: [0.5, 0.7, 0.9]}).rename_axis(index="param")


def mean(*args):
    return sum(args) / len(args)


transformed = (
    cb.Graph(graph.to_networkx())
    .map(df)
    .map(params)
    .reduce(Normalized, name='Normalized', attrs={'reduce': mean}, index='bank')
    .reduce('Normalized', name='Normalized2', attrs={'reduce': mean}, index='param')
)
transformed
graph2 = sciline.task_graph.TaskGraph.from_networkx(transformed.graph)
graph2.visualize(graph_attr={"rankdir": "LR"})
graph2.compute('Normalized2')