// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module fractions

import math.big

// Fraction stores a numerator and denominator using the existing `i64` backend.
pub struct Fraction {
pub:
	n          i64
	d          i64
	is_reduced bool
}

// Rational stores a numerator and denominator using the backend type `T`.
//
// Supported backend types are signed builtin integers and `math.big.Integer`.
pub struct Rational[T] {
pub:
	n          T
	d          T
	is_reduced bool
}

// fraction creates an `i64`-backed Fraction.
//
// The denominator must be non-zero. Negative denominators are normalized so
// that the stored denominator is always positive.
pub fn fraction(n i64, d i64) Fraction {
	return rational_to_fraction(rational[i64](n, d))
}

// big_fraction creates a `math.big.Integer`-backed Rational.
//
// The denominator must be non-zero. Negative denominators are normalized so
// that the stored denominator is always positive.
pub fn big_fraction(n big.Integer, d big.Integer) Rational[big.Integer] {
	return rational[big.Integer](n, d)
}

// rational creates a Rational backed by the concrete numeric type `T`.
//
// Fractions created this way are not reduced automatically, but `is_reduced`
// reflects whether the input was already in lowest terms.
pub fn rational[T](n T, d T) Rational[T] {
	ensure_supported_backend[T]()
	if is_zero_value[T](d) {
		panic('Denominator cannot be zero')
	}
	if is_negative_value[T](d) {
		return rational[T](negate_value[T](n), negate_value[T](d))
	}
	return Rational[T]{
		n:          n
		d:          d
		is_reduced: gcd_values[T](n, d) == one_value[T]()
	}
}

// str returns the fraction in `n/d` form.
pub fn (f Fraction) str() string {
	return f.to_rational().str()
}

// Fraction add using operator overloading.
pub fn (f1 Fraction) + (f2 Fraction) Fraction {
	return rational_to_fraction(f1.to_rational() + f2.to_rational())
}

// Fraction subtract using operator overloading.
pub fn (f1 Fraction) - (f2 Fraction) Fraction {
	return rational_to_fraction(f1.to_rational() - f2.to_rational())
}

// Fraction multiply using operator overloading.
pub fn (f1 Fraction) * (f2 Fraction) Fraction {
	return rational_to_fraction(f1.to_rational() * f2.to_rational())
}

// Fraction divide using operator overloading.
pub fn (f1 Fraction) / (f2 Fraction) Fraction {
	return rational_to_fraction(f1.to_rational() / f2.to_rational())
}

// negate returns the additive inverse of the Fraction.
pub fn (f Fraction) negate() Fraction {
	return rational_to_fraction(f.to_rational().negate())
}

// reciprocal returns the reciprocal of the Fraction.
pub fn (f Fraction) reciprocal() Fraction {
	return rational_to_fraction(f.to_rational().reciprocal())
}

// reduce returns the Fraction reduced to lowest terms.
pub fn (f Fraction) reduce() Fraction {
	return rational_to_fraction(f.to_rational().reduce())
}

// f64 converts the Fraction to 64-bit floating point.
pub fn (f Fraction) f64() f64 {
	return f.to_rational().f64()
}

// return true if f1 == f2.
pub fn (f1 Fraction) == (f2 Fraction) bool {
	return cmp_fraction(f1, f2) == 0
}

// return true if f1 < f2.
pub fn (f1 Fraction) < (f2 Fraction) bool {
	return cmp_fraction(f1, f2) < 0
}

// str returns the fraction in `n/d` form.
pub fn (f Rational[T]) str() string {
	return '${f.n}/${f.d}'
}

//
// + ---------------------+
// | Arithmetic functions.|
// + ---------------------+
//
// These are implemented from Knuth, TAOCP Vol 2. Section 4.5
//
// Returns a correctly reduced result for both addition and subtraction.
// NOTE: requires reduced inputs.
fn general_addition_result[T](f1 Rational[T], f2 Rational[T], addition bool) Rational[T] {
	d1 := gcd_values[T](f1.d, f2.d)
	// d1 happens to be 1 around 600/(pi)^2 or 61 percent of the time (Theorem 4.5.2D)
	if d1 == one_value[T]() {
		num1n2d := f1.n * f2.d
		num1d2n := f1.d * f2.n
		if addition {
			return Rational[T]{
				n:          add_values[T](num1n2d, num1d2n)
				d:          f1.d * f2.d
				is_reduced: true
			}
		}
		return Rational[T]{
			n:          sub_values[T](num1n2d, num1d2n)
			d:          f1.d * f2.d
			is_reduced: true
		}
	}
	// Here d1 > 1.
	f1den := f1.d / d1
	f2den := f2.d / d1
	term1 := f1.n * f2den
	term2 := f2.n * f1den
	if addition {
		t := add_values[T](term1, term2)
		d2 := gcd_values[T](t, d1)
		return Rational[T]{
			n:          t / d2
			d:          f1den * (f2.d / d2)
			is_reduced: true
		}
	}
	t := sub_values[T](term1, term2)
	d2 := gcd_values[T](t, d1)
	return Rational[T]{
		n:          t / d2
		d:          f1den * (f2.d / d2)
		is_reduced: true
	}
}

// Fraction add using operator overloading.
pub fn (f1 Rational[T]) + (f2 Rational[T]) Rational[T] {
	return general_addition_result[T](f1.reduce(), f2.reduce(), true)
}

// Fraction subtract using operator overloading.
pub fn (f1 Rational[T]) - (f2 Rational[T]) Rational[T] {
	return general_addition_result[T](f1.reduce(), f2.reduce(), false)
}

// Returns a correctly reduced result for both multiplication and division.
// NOTE: requires reduced inputs.
fn general_multiplication_result[T](f1 Rational[T], f2 Rational[T], multiplication bool) Rational[T] {
	// * Theorem: If f1 and f2 are reduced i.e. gcd(f1.n, f1.d) ==  1 and gcd(f2.n, f2.d) == 1,
	// then gcd(f1.n * f2.n, f1.d * f2.d) == gcd(f1.n, f2.d) * gcd(f1.d, f2.n)
	// * Knuth poses this an exercise for 4.5.1. - Exercise 2
	// * Also, note that:
	// The terms are flipped for multiplication and division, so the gcds must be calculated carefully
	// We do multiple divisions in order to prevent any possible overflows.
	// * One more thing:
	// if d = gcd(a, b) for example, then d divides both a and b
	if multiplication {
		d1 := gcd_values[T](f1.n, f2.d)
		d2 := gcd_values[T](f1.d, f2.n)
		return Rational[T]{
			n:          (f1.n / d1) * (f2.n / d2)
			d:          (f2.d / d1) * (f1.d / d2)
			is_reduced: true
		}
	} else {
		d1 := gcd_values[T](f1.n, f2.n)
		d2 := gcd_values[T](f1.d, f2.d)
		return Rational[T]{
			n:          (f1.n / d1) * (f2.d / d2)
			d:          (f2.n / d1) * (f1.d / d2)
			is_reduced: true
		}
	}
}

// Fraction multiply using operator overloading.
pub fn (f1 Rational[T]) * (f2 Rational[T]) Rational[T] {
	return general_multiplication_result[T](f1.reduce(), f2.reduce(), true)
}

// Fraction divide using operator overloading.
pub fn (f1 Rational[T]) / (f2 Rational[T]) Rational[T] {
	if is_zero_value[T](f2.n) {
		panic('Cannot divide by zero')
	}
	// If the second fraction is negative, it will mess up the sign.
	// We need a positive denominator.
	if is_negative_value[T](f2.n) {
		return f1.negate() / f2.negate()
	}
	return general_multiplication_result[T](f1.reduce(), f2.reduce(), false)
}

// negate returns the additive inverse of the Rational.
pub fn (f Rational[T]) negate() Rational[T] {
	return Rational[T]{
		n:          negate_value[T](f.n)
		d:          f.d
		is_reduced: f.is_reduced
	}
}

// reciprocal returns the reciprocal of the Rational.
pub fn (f Rational[T]) reciprocal() Rational[T] {
	if is_zero_value[T](f.n) {
		panic('Denominator cannot be zero')
	}
	return rational[T](f.d, f.n)
}

// reduce returns the Rational reduced to lowest terms.
pub fn (f Rational[T]) reduce() Rational[T] {
	if f.is_reduced {
		return f
	}
	cf := gcd_values[T](f.n, f.d)
	return Rational[T]{
		n:          f.n / cf
		d:          f.d / cf
		is_reduced: true
	}
}

// f64 converts the Rational to 64-bit floating point.
pub fn (f Rational[T]) f64() f64 {
	return to_f64_value[T](f.n) / to_f64_value[T](f.d)
}

//
// + ------------------+
// | Utility functions.|
// + ------------------+
//

fn ensure_supported_backend[T]() {
	$if T is big.Integer || T is i8 || T is i16 || T is i32 || T is i64 || T is int {
	} $else {
		$compile_error('fractions: Rational only supports signed builtin integers and math.big.Integer backends')
	}
}

@[inline]
fn (f Fraction) to_rational() Rational[i64] {
	return Rational[i64]{
		n:          f.n
		d:          f.d
		is_reduced: f.is_reduced
	}
}

@[inline]
fn rational_to_fraction(r Rational[i64]) Fraction {
	return Fraction{
		n:          r.n
		d:          r.d
		is_reduced: r.is_reduced
	}
}

fn cmp_fraction(f1 Fraction, f2 Fraction) int {
	return cmp_values[big.Integer](to_big_integer[i64](f1.n) * to_big_integer[i64](f2.d),
		to_big_integer[i64](f2.n) * to_big_integer[i64](f1.d))
}

@[inline]
fn zero_value[T]() T {
	$if T is big.Integer {
		return big.zero_int
	} $else {
		return T(0)
	}
}

@[inline]
fn one_value[T]() T {
	$if T is big.Integer {
		return big.one_int
	} $else {
		return T(1)
	}
}

@[inline]
fn is_zero_value[T](value T) bool {
	return value == zero_value[T]()
}

@[inline]
fn is_negative_value[T](value T) bool {
	$if T is big.Integer {
		return value < big.zero_int
	} $else {
		return value < 0
	}
}

@[inline]
fn negate_value[T](value T) T {
	$if T is big.Integer {
		return value.neg()
	} $else {
		return -value
	}
}

@[inline]
fn abs_value[T](value T) T {
	$if T is big.Integer {
		return value.abs()
	} $else {
		return if value < 0 { -value } else { value }
	}
}

fn gcd_values[T](a T, b T) T {
	$if T is big.Integer {
		return a.gcd(b)
	} $else {
		mut x := abs_value[T](a)
		mut y := abs_value[T](b)
		for y != zero_value[T]() {
			x %= y
			if x == zero_value[T]() {
				return y
			}
			y %= x
		}
		return x
	}
}

@[inline]
fn to_f64_value[T](value T) f64 {
	$if T is big.Integer {
		return value.str().f64()
	} $else {
		return f64(value)
	}
}

@[inline]
fn to_big_integer[T](value T) big.Integer {
	$if T is big.Integer {
		return value
	} $else {
		return big.integer_from_i64(i64(value))
	}
}

@[inline]
fn add_values[T](a T, b T) T {
	return a + b
}

@[inline]
fn sub_values[T](a T, b T) T {
	return a - b
}

// cmp compares the two arguments, returns 0 when equal, 1 when the first is bigger, -1 otherwise.
fn cmp[T](f1 Rational[T], f2 Rational[T]) int {
	$if T is big.Integer {
		return cmp_values[T](f1.n * f2.d, f2.n * f1.d)
	} $else {
		return cmp_values[big.Integer](to_big_integer[T](f1.n) * to_big_integer[T](f2.d),
			to_big_integer[T](f2.n) * to_big_integer[T](f1.d))
	}
}

fn cmp_values[T](a T, b T) int {
	if a == b {
		return 0
	}
	return if a > b { 1 } else { -1 }
}

// return true if f1 == f2.
pub fn (f1 Rational[T]) == (f2 Rational[T]) bool {
	return cmp[T](f1, f2) == 0
}

// return true if f1 < f2.
pub fn (f1 Rational[T]) < (f2 Rational[T]) bool {
	return cmp[T](f1, f2) < 0
}