# Generates code containing quadrature rules for the Gauss-Lobatto
from quadpy import c1
from pathlib import Path
import sys

# Include all lower-order rules up to a maximum value, then power of two after that
LOWER_ORDER_MAX = 32
higher_order_set = [2**i for i in range(6, 10)]
n_set = list(range(2, LOWER_ORDER_MAX + 1)) + higher_order_set

file_template = r"""//! Code generated by %filename%.
//! Do not change this file directly. Instead, modify the codegen script.
use crate::Rule;

%functions%

// Documentation for this function is provided in the outer (non-codegen) function.
pub fn try_gauss_lobatto(n: usize) -> Option<Rule<1>> {
    match n {
%cases%
        _ => None,
    }
}
"""

fn_template = r"""#[rustfmt::skip]
pub fn gauss_lobatto_%n%() -> Rule<1> {
    let weights = vec![
%weights%
    ];
    let points = vec![
%points%
    ];
    (weights, points)
}"""


def build_gauss_lobatto_code(n):
    scheme = c1.gauss_lobatto(n)
    indent = " " * 8
    weights_str = "\n".join([indent + str(w) + "," for w in scheme.weights])
    points_str = "\n".join([indent + "[" + str(xi) + "]," for xi in scheme.points])
    fn = fn_template.replace('%n%', str(n)).replace('%weights%', weights_str).replace('%points%', points_str)
    return fn


def build_gauss_lobatto_case(n):
    indent = " " * 8
    return f"{indent}{n} => Some(gauss_lobatto_{n}()),"


function_codes = "\n\n".join([build_gauss_lobatto_code(n) for n in n_set])
cases = "\n".join([build_gauss_lobatto_case(n) for n in n_set])
script_filename = Path(__file__).name
file_str = file_template \
    .replace('%filename%', "tools/" + script_filename) \
    .replace('%functions%', function_codes) \
    .replace('%cases%', cases)

codegen_dir_name = "src/codegen"
codegen_dir = Path(__file__).parent.parent.joinpath(codegen_dir_name)

# Generate a doc stub that will be included in the manual docs (generated so things don't get out of sync).
higher_order_set_str = ", ".join([str(n) for n in higher_order_set])
doc_stub = f"""Rules are provided for all lower-order rules up to `n <= {LOWER_ORDER_MAX}`,
plus `n = ..., {higher_order_set_str}`.
"""

if codegen_dir.exists():
    filepath = codegen_dir.joinpath("gauss_lobatto_rules.rs")
    print(f"Writing generated code to {filepath}...")
    with open(filepath, 'w') as f:
        f.write(file_str)
    print("Successfully wrote generated code.")
    docpath = codegen_dir.joinpath("gauss_lobatto_rules_doc_stub.md")
    print(f"Writing doc stub to {docpath}...")
    with open(docpath, 'w') as f:
        f.write(doc_stub)
    print(f"Successfully wrote doc stub.")
else:
    print(f"{codegen_dir_name} does not exist")
    sys.exit(1)