pub struct AtmosphericFlow {
    pub name: &'static str,
    pub lengthkm: f64,
    pub flowratem3s: f64,
    pub zonalwidthkm: f64,
    pub slope: f64,
}

impl AtmosphericFlow {
    pub fn manningvelocity(&self, hydraulicradiusm: f64, roughnessn: f64) -> f64 {
        if roughnessn.abs() < f64::EPSILON {
            return 0.0;
        }
        (1.0 / roughnessn) * hydraulicradiusm.powf(2.0 / 3.0) * self.slope.sqrt()
    }

    pub fn froudenumber(&self, velocityms: f64, depthm: f64) -> f64 {
        if depthm.abs() < 1e-10 {
            return 0.0;
        }
        velocityms / (*crate::SURFACEGRAVITY * depthm.abs()).sqrt()
    }

    pub fn reynoldsnumber(&self, velocityms: f64, depthm: f64) -> f64 {
        let kinematicviscosity = 1e-5;
        velocityms * depthm / kinematicviscosity
    }

    pub fn streampowerwperm(&self) -> f64 {
        *crate::ONEBARDENSITY * *crate::SURFACEGRAVITY * self.flowratem3s * self.slope
    }

    pub fn specificdischargems(&self) -> f64 {
        if self.zonalwidthkm.abs() < 1e-30 {
            return 0.0;
        }
        self.flowratem3s / (self.zonalwidthkm * 1e3)
    }
}

pub fn equatorialjet() -> AtmosphericFlow {
    AtmosphericFlow {
        name: "Equatorial Retrograde Jet",
        lengthkm: 77000.0,
        flowratem3s: 1e12,
        zonalwidthkm: 10000.0,
        slope: 1e-8,
    }
}

pub fn polarjet() -> AtmosphericFlow {
    AtmosphericFlow {
        name: "Polar Prograde Jet",
        lengthkm: 20000.0,
        flowratem3s: 5e11,
        zonalwidthkm: 5000.0,
        slope: 1e-8,
    }
}

pub struct FlowSegment {
    pub flow: AtmosphericFlow,
    pub downstreamidx: Option<usize>,
}

pub struct FlowNetwork {
    pub segments: Vec<FlowSegment>,
}

impl Default for FlowNetwork {
    fn default() -> Self {
        Self::new()
    }
}

impl FlowNetwork {
    pub fn new() -> Self {
        Self {
            segments: Vec::new(),
        }
    }

    pub fn addheadwater(&mut self, flow: AtmosphericFlow) -> usize {
        let idx = self.segments.len();
        self.segments.push(FlowSegment {
            flow,
            downstreamidx: None,
        });
        idx
    }

    pub fn connect(&mut self, tributaryidx: usize, mainidx: usize) {
        self.segments[tributaryidx].downstreamidx = Some(mainidx);
    }

    pub fn routedischarge(&self) -> Vec<f64> {
        let n = self.segments.len();
        let mut discharge = vec![0.0f64; n];
        for (i, seg) in self.segments.iter().enumerate() {
            discharge[i] = seg.flow.flowratem3s;
        }
        let mut indegree = vec![0u32; n];
        for seg in &self.segments {
            if let Some(ds) = seg.downstreamidx {
                indegree[ds] += 1;
            }
        }
        let mut queue: Vec<usize> = (0..n).filter(|&i| indegree[i] == 0).collect();
        let mut order = Vec::with_capacity(n);
        let mut remainingin = indegree.clone();
        while let Some(idx) = queue.pop() {
            order.push(idx);
            if let Some(ds) = self.segments[idx].downstreamidx {
                remainingin[ds] -= 1;
                if remainingin[ds] == 0 {
                    queue.push(ds);
                }
            }
        }
        for &idx in &order {
            if let Some(ds) = self.segments[idx].downstreamidx {
                discharge[ds] += discharge[idx];
            }
        }
        discharge
    }

    pub fn outletdischarge(&self) -> f64 {
        let discharge = self.routedischarge();
        self.segments
            .iter()
            .enumerate()
            .filter(|(_, seg)| seg.downstreamidx.is_none())
            .map(|(i, _)| discharge[i])
            .fold(0.0f64, |a, b| a.max(b))
    }
}

pub fn neptunejetnetwork() -> FlowNetwork {
    let mut net = FlowNetwork::new();
    let eq = net.addheadwater(equatorialjet());
    let pol = net.addheadwater(polarjet());
    net.connect(pol, eq);
    net
}
