use std::marker::PhantomData;

pub trait Rev {
    type Input;
    type Output;

    fn forward(&self, i: Self::Input) -> Self::Output;
    fn reverse(&self, o: Self::Output) -> Self::Input;
}

#[derive(Clone)]
pub struct Addition<A> {
    amount: A
}
impl<A> Rev for Addition<A>
where
    A: std::ops::Add<Output = A> + std::ops::Sub<Output = A> + Copy
{
    type Input = A;
    type Output = A;

    fn forward(&self, a: A) -> A {
        a + self.amount
    }

    fn reverse(&self, b: A) -> A {
        b - self.amount
    }
}
pub fn plus<A>(amount: A) -> Addition<A> {
    Addition { amount }
}

// more blocks

#[derive(Clone)]
pub struct Negate<A>(PhantomData<A>);
impl<A: std::ops::Neg<Output = A>> Rev for Negate<A> {
    type Input = A;
    type Output = A;

    fn forward(&self, a: Self::Input) -> Self::Output {
        -a
    }

    fn reverse(&self, b: Self::Output) -> Self::Input {
        -b
    }
}
pub fn neg<A>() -> Negate<A> {
    Negate(PhantomData)
}

#[derive(Clone)]
pub struct Multiplication<A> {
    amount: A
}
impl<A> Rev for Multiplication<A>
where
    A: std::ops::Mul<Output = A> + std::ops::Div<Output = A> + Copy,
{
    type Input = A;
    type Output = A;

    fn forward(&self, a: A) -> A {
        a * self.amount
    }

    fn reverse(&self, b: A) -> A {
        b / self.amount
    }
}
pub fn times<A>(amount: A) -> Multiplication<A> {
    Multiplication { amount }
}

#[derive(Clone)]
pub struct DiscreteMap<A, B>(Vec<(A, B)>);
impl<A: PartialEq + Clone, B: PartialEq + Clone> Rev for DiscreteMap<A, B> {
    type Input = A;
    type Output = B;
    
    fn forward(&self, a: A) -> B {
        for i in self.0.iter() {
            if i.0 == a {
                return i.1.clone()
            }
        }
        panic!()
    }

    fn reverse(&self, b: B) -> A {
        for i in self.0.iter() {
            if i.1 == b {
                return i.0.clone()
            }
        }
        panic!()
    }
}
pub fn discrete<A, B>(vec: Vec<(A, B)>) -> DiscreteMap<A, B> {
    DiscreteMap(vec)
}

// combinators

#[derive(Clone)]
pub struct Pipe<X, Y>(X, Y);
impl<X, Y> Rev for Pipe<X, Y>
where
    X: Rev,
    Y: Rev<Input = X::Output>
{
    type Input = X::Input;
    type Output = Y::Output;

    fn forward(&self, a: Self::Input) -> Self::Output {
        self.1.forward(self.0.forward(a))
    }

    fn reverse(&self, c: Self::Output) -> Self::Input {
        self.0.reverse(self.1.reverse(c))
    }
}

#[derive(Clone)]
pub struct Inverse<X>(X);
impl<X: Rev> Rev for Inverse<X> {
    type Input = X::Output;
    type Output = X::Input;

    fn forward(&self, a: Self::Input) -> Self::Output {
        self.0.reverse(a)
    }

    fn reverse(&self, b: Self::Output) -> Self::Input {
        self.0.forward(b)
    }
}

pub trait RevExt: Rev + Sized {
    fn then<N: Rev<Input = Self::Output>>(self, other: N) -> Pipe<Self, N> {
        Pipe(self, other)
    }

    fn inverse(self) -> Inverse<Self> {
        Inverse(self)
    }
}
impl<T: Rev> RevExt for T {}

pub fn under<Op, Val>(op: Op, val: Val) -> impl Rev<Input = Op::Input, Output = Op::Input>
where
    Op: Rev + Clone,
    Val: Rev<Input = Op::Output, Output = Op::Output>
{
    op.clone().then(val).then(op.inverse())
}