//! river2d_core
//!
//! FastAPIから標準入力JSONで呼び出す、簡易2次元河川不定流計算エンジンです。
//!
//! 注意:
//! - 実務用の高精度ソルバーではなく、Web API構成を説明するためのデモです。
//! - 計算は浅水流の厳密なRiemannソルバーではなく、水面勾配とManning式から
//!   セル間フラックスを求める拡散波近似に近い陽解法です。
//! - 実務ではCFL条件、境界条件、ドライ/ウェット処理、地形補正、観測値検証を追加してください。

use serde::Serialize;
use serde_json::{json, Value};
use std::io::{self, Read};

#[derive(Debug, Clone, Serialize)]
struct HydroPoint {
    time_s: f64,
    inflow_m3s: f64,
    outflow_m3s: f64,
    max_depth_m: f64,
    flooded_area_m2: f64,
}

fn get_f64(v: &Value, key: &str, default: f64) -> f64 {
    v.get(key).and_then(Value::as_f64).unwrap_or(default)
}

fn get_usize(v: &Value, key: &str, default: usize) -> usize {
    v.get(key)
        .and_then(Value::as_u64)
        .map(|x| x as usize)
        .unwrap_or(default)
}

fn get_bool(v: &Value, key: &str, default: bool) -> bool {
    v.get(key).and_then(Value::as_bool).unwrap_or(default)
}

fn idx(i: usize, j: usize, ny: usize) -> usize {
    i * ny + j
}

fn round3(x: f64) -> f64 {
    (x * 1000.0).round() / 1000.0
}

fn round4(x: f64) -> f64 {
    (x * 10000.0).round() / 10000.0
}

fn inflow_at(t: f64, base: f64, peak: f64, peak_time: f64) -> f64 {
    let sigma = (peak_time * 0.305).max(10.0);
    base + peak * (-0.5 * ((t - peak_time) / sigma).powi(2)).exp()
}

fn to_rows(values: &[f64], nx: usize, ny: usize) -> Vec<Vec<f64>> {
    let mut rows = Vec::with_capacity(ny);
    for j in 0..ny {
        let mut row = Vec::with_capacity(nx);
        for i in 0..nx {
            row.push(round4(values[idx(i, j, ny)]));
        }
        rows.push(row);
    }
    rows
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut input = String::new();
    io::stdin().read_to_string(&mut input)?;
    let req: Value = serde_json::from_str(&input)?;

    let nx = get_usize(&req, "nx", 72).clamp(20, 220);
    let ny = get_usize(&req, "ny", 24).clamp(8, 120);
    let dx = get_f64(&req, "dx", 20.0).max(0.1);
    let dy = get_f64(&req, "dy", 20.0).max(0.1);
    let dt = get_f64(&req, "dt", 1.0).clamp(0.01, 10.0);
    let steps = get_usize(&req, "steps", 2400).clamp(1, 7200);
    let manning_n = get_f64(&req, "manning_n", 0.035).clamp(0.005, 0.20);
    let bed_slope = get_f64(&req, "bed_slope", 0.005).clamp(0.000001, 0.02);
    let initial_depth = get_f64(&req, "initial_depth", 0.3).clamp(0.01, 10.0);
    let inflow_base = get_f64(&req, "inflow_base_m3s", 55.0).max(0.0);
    let inflow_peak = get_f64(&req, "inflow_peak_m3s", 230.0).max(0.0);
    let peak_time = get_f64(&req, "peak_time_s", 360.0).max(1.0);
    let velocity_limit = get_f64(&req, "velocity_limit_ms", 6.0).clamp(0.1, 20.0);
    let output_interval = get_usize(&req, "output_interval", 10).clamp(1, 600);

    let null_value = Value::Null;
    let breach = req.get("breach").unwrap_or(&null_value);
    let breach_enabled = get_bool(breach, "enabled", true);
    let breach_x_start = get_usize(breach, "x_start", nx * 3 / 5).min(nx.saturating_sub(1));
    let breach_x_end = get_usize(breach, "x_end", breach_x_start + 5).min(nx.saturating_sub(1));
    let breach_right_bank = get_bool(breach, "right_bank", true);

    let area = dx * dy;
    let y_center = (ny as f64 - 1.0) / 2.0;
    let channel_half_cells = 3.2_f64;
    let cell_count = nx * ny;

    let mut z = vec![0.0; cell_count];
    let mut h = vec![0.0; cell_count];
    let mut max_h = vec![0.0; cell_count];

    for i in 0..nx {
        let longitudinal = (nx - 1 - i) as f64 * dx * bed_slope;
        for j in 0..ny {
            let jj = j as f64;
            let in_channel = (jj - y_center).abs() <= channel_half_cells;
            let mut base = if in_channel {
                -1.2 + 0.018 * (jj - y_center).abs().powi(2)
            } else {
                0.65 + 0.06 * ((jj - y_center).abs() - channel_half_cells)
            };

            if !in_channel {
                let near_bank = ((jj - y_center).abs() - (channel_half_cells + 1.0)).abs() < 1.2;
                let right_side = jj > y_center;
                let breach_side = if breach_right_bank { right_side } else { !right_side };
                let in_breach_x = breach_enabled && i >= breach_x_start && i <= breach_x_end;
                if near_bank && !(in_breach_x && breach_side) {
                    base += 0.55;
                } else if near_bank && in_breach_x && breach_side {
                    base -= 0.20;
                }
            }

            let k = idx(i, j, ny);
            z[k] = longitudinal + base;
            if in_channel {
                h[k] = initial_depth;
            }
            max_h[k] = h[k];
        }
    }

    let mut hydrograph: Vec<HydroPoint> = Vec::new();
    let mut last_qx = vec![0.0; (nx + 1) * ny];
    let mut last_qy = vec![0.0; nx * (ny + 1)];

    for step in 0..=steps {
        let t = step as f64 * dt;
        let mut eta = vec![0.0; cell_count];
        for k in 0..cell_count {
            eta[k] = z[k] + h[k];
        }

        let mut qx = vec![0.0; (nx + 1) * ny]; // 正: +x方向
        let mut qy = vec![0.0; nx * (ny + 1)]; // 正: +y方向

        // x方向の内部フラックス
        for i in 1..nx {
            for j in 0..ny {
                let l = idx(i - 1, j, ny);
                let r = idx(i, j, ny);
                let slope = (eta[l] - eta[r]) / dx;
                let hf = 0.5 * (h[l] + h[r]);
                if hf > 1.0e-4 {
                    let raw_q = (1.0 / manning_n)
                        * hf.powf(5.0 / 3.0)
                        * (slope.abs() + 1.0e-12).sqrt()
                        * slope.signum()
                        * dy;
                    let cap = velocity_limit * hf * dy;
                    qx[i * ny + j] = raw_q.clamp(-cap, cap);
                }
            }
        }

        // 下流端の自由流出
        for j in 0..ny {
            let k = idx(nx - 1, j, ny);
            if h[k] > 1.0e-4 {
                let raw_q = (1.0 / manning_n) * h[k].powf(5.0 / 3.0) * bed_slope.sqrt() * dy;
                qx[nx * ny + j] = raw_q.min(velocity_limit * h[k] * dy);
            }
        }

        // y方向の内部フラックス
        for i in 0..nx {
            for j in 1..ny {
                let b = idx(i, j - 1, ny);
                let tcell = idx(i, j, ny);
                let slope = (eta[b] - eta[tcell]) / dy;
                let hf = 0.5 * (h[b] + h[tcell]);
                if hf > 1.0e-4 {
                    let raw_q = (1.0 / manning_n)
                        * hf.powf(5.0 / 3.0)
                        * (slope.abs() + 1.0e-12).sqrt()
                        * slope.signum()
                        * dx;
                    let cap = velocity_limit * hf * dx;
                    qy[i * (ny + 1) + j] = raw_q.clamp(-cap, cap);
                }
            }
        }

        // 陽解法で水深を更新するため、まず各セルの流出量を集計し、
        // 1ステップでセル内水量を超えて流出しないように制限します。
        let mut outflow = vec![0.0; cell_count];

        for i in 1..nx {
            for j in 0..ny {
                let q = qx[i * ny + j];
                if q > 0.0 {
                    outflow[idx(i - 1, j, ny)] += q;
                } else {
                    outflow[idx(i, j, ny)] += -q;
                }
            }
        }
        for j in 0..ny {
            outflow[idx(nx - 1, j, ny)] += qx[nx * ny + j].max(0.0);
        }
        for i in 0..nx {
            for j in 1..ny {
                let q = qy[i * (ny + 1) + j];
                if q > 0.0 {
                    outflow[idx(i, j - 1, ny)] += q;
                } else {
                    outflow[idx(i, j, ny)] += -q;
                }
            }
        }

        let mut scale = vec![1.0; cell_count];
        for k in 0..cell_count {
            if outflow[k] > 1.0e-12 {
                scale[k] = ((h[k] * area) / (dt * outflow[k])).min(1.0);
            }
        }

        for i in 1..nx {
            for j in 0..ny {
                let q = qx[i * ny + j];
                if q > 0.0 {
                    qx[i * ny + j] *= scale[idx(i - 1, j, ny)];
                } else {
                    qx[i * ny + j] *= scale[idx(i, j, ny)];
                }
            }
        }
        for j in 0..ny {
            qx[nx * ny + j] *= scale[idx(nx - 1, j, ny)];
        }
        for i in 0..nx {
            for j in 1..ny {
                let q = qy[i * (ny + 1) + j];
                if q > 0.0 {
                    qy[i * (ny + 1) + j] *= scale[idx(i, j - 1, ny)];
                } else {
                    qy[i * (ny + 1) + j] *= scale[idx(i, j, ny)];
                }
            }
        }

        if step % output_interval == 0 {
            let out_q: f64 = (0..ny).map(|j| qx[nx * ny + j].max(0.0)).sum();
            let max_depth_now = h.iter().fold(0.0_f64, |a, b| a.max(*b));
            let flooded_area = h.iter().filter(|v| **v > 0.05).count() as f64 * area;
            hydrograph.push(HydroPoint {
                time_s: round3(t),
                inflow_m3s: round3(inflow_at(t, inflow_base, inflow_peak, peak_time)),
                outflow_m3s: round3(out_q),
                max_depth_m: round4(max_depth_now),
                flooded_area_m2: round3(flooded_area),
            });
        }

        if step == steps {
            last_qx = qx;
            last_qy = qy;
            break;
        }

        let mut next_h = h.clone();
        for i in 0..nx {
            for j in 0..ny {
                let k = idx(i, j, ny);
                let dh = (
                    qx[i * ny + j]
                    - qx[(i + 1) * ny + j]
                    + qy[i * (ny + 1) + j]
                    - qy[i * (ny + 1) + j + 1]
                ) * dt / area;
                next_h[k] = (h[k] + dh).max(0.0);
            }
        }

        // 上流端の流入ハイドログラフを中央河道セルへ分配します。
        let q_in = inflow_at(t, inflow_base, inflow_peak, peak_time);
        let channel_js: Vec<usize> = (0..ny)
            .filter(|j| (*j as f64 - y_center).abs() <= channel_half_cells)
            .collect();
        let add_depth = q_in * dt / (channel_js.len() as f64 * area);
        for j in channel_js {
            next_h[idx(0, j, ny)] += add_depth;
        }

        h = next_h;
        for k in 0..cell_count {
            if h[k] > max_h[k] {
                max_h[k] = h[k];
            }
        }
    }

    let mut velocity = vec![0.0; cell_count];
    for i in 0..nx {
        for j in 0..ny {
            let k = idx(i, j, ny);
            let hh = h[k].max(0.05);
            let u = 0.5 * (last_qx[i * ny + j] + last_qx[(i + 1) * ny + j]) / (dy * hh);
            let v = 0.5 * (last_qy[i * (ny + 1) + j] + last_qy[i * (ny + 1) + j + 1]) / (dx * hh);
            let speed = (u * u + v * v).sqrt();
            velocity[k] = speed.min(velocity_limit);
        }
    }

    let peak_inflow = hydrograph.iter().fold(0.0_f64, |a, p| a.max(p.inflow_m3s));
    let peak_outflow = hydrograph.iter().fold(0.0_f64, |a, p| a.max(p.outflow_m3s));
    let max_depth = max_h.iter().fold(0.0_f64, |a, b| a.max(*b));
    let max_velocity = velocity.iter().fold(0.0_f64, |a, b| a.max(*b));
    let flooded_area_final = h.iter().filter(|v| **v > 0.05).count() as f64 * area;

    let response = json!({
        "metadata": {
            "model": "fastapi-rust-river2d-demo",
            "solver": "Rust CLI / simplified 2D unsteady shallow-water diffusion wave",
            "nx": nx,
            "ny": ny,
            "dx": dx,
            "dy": dy,
            "dt": dt,
            "steps": steps,
            "peak_inflow_m3s": round3(peak_inflow),
            "peak_outflow_m3s": round3(peak_outflow),
            "max_depth_m": round4(max_depth),
            "max_velocity_ms": round4(max_velocity),
            "flooded_area_final_m2": round3(flooded_area_final),
            "note": "デモ用の簡易モデルです。実務ではCFL条件、境界条件、粗度、地形処理、観測値検証を追加してください。"
        },
        "hydrograph": hydrograph,
        "maps": {
            "bed_elevation_m": to_rows(&z, nx, ny),
            "final_depth_m": to_rows(&h, nx, ny),
            "max_depth_m": to_rows(&max_h, nx, ny),
            "final_velocity_ms": to_rows(&velocity, nx, ny)
        }
    });

    println!("{}", serde_json::to_string_pretty(&response)?);
    Ok(())
}