:- module(
  tropical,
  [
    tropicalize/0
    % grammar
    %Public API
    ]
  ).

% Only for separate compilation/linting
:- use_module(doc).
:- use_module(biocham).
:- use_module(library(clpfd)).


:- doc('Tropical algebra can be used to reason about orders of magnitude of molecular concentrations, kinetic parameters and reactions rates.
The solutions to tropical equilibration problems provide candidates for regimes exhibiting fast-slow dynamics leading to model reductions \\cite{SFR14amb}.').


% store monomials of -inf degree
:- dynamic(is_zero/1).

% equilibrate(?LPos, ?LNeg)
%
% To equilibrate, the positive term of min degree should be of same degree as
% the neg term of min degree
equilibrate([], []) :-
  !.

equilibrate(LPos, LNeg) :-
  get_min(LPos, Min),
  get_min(LNeg, Min).


% equilibrate_rec(-L)
%
% call equilibrate for each pair (Pos, Neg) of list L
equilibrate_rec([]).

equilibrate_rec([(Pos, Neg) | E]) :-
  equilibrate(Pos, Neg),
  equilibrate_rec(E).


% get_min(?L, ?M)
%
% set M as smaller than all variables of L
% C is the number of variables of L actually equal to M
%
% since here we use C #>=1, we rely on swipl's min fd expr
% get_min([], +inf).
get_min([A], M) :-
  !,
  M #= A.
get_min([H | T], M) :-
  M #= min(H, MM),
  get_min(T, MM).

% get_min(L, M) :-
%    get_min_and_count(L, M, C),
%    % here we enforce at least 2-term equilibration (i.e. max of positive is
%    % equal to max of neg and all others negligible)
%    % for strict 2-term equilibration one would have C #= 1.
%    C #>= 1.


% get_min_and_count(?L, ?M, ?C)
%
% sets M as smaller than all variables of the list given as first arg
% sets C to the count of variables of that list actually equal to M
get_min_and_count([H | T], M, C) :-
  M #=< H,
  B #<==> M #= H,
  C #= B + CC,
  get_min_and_count(T, M, CC).

get_min_and_count([], _, 0).


% conserve(?L, -D)
%
% a conservation law implies that when epsilon goes to infinity, the remaining
% terms (those of minimal degree) should be of degree DD (since the sum is
% equal to a constant D) and at least two (since otherwise it would force one
% variable to be always constant)
conserve(_, 0) :-
  % FIXME no idea what to do :-/
  !.

conserve(L, D) :-
  get_degree(D, DD),
  % we do not use get_min here, even with C #>= 1 since L is a long list and
  % deeply nested min is not that efficient...
  get_min_and_count(L, DD, C),
  C #>= 1.
% C #>= 2.     % oscillations FIXME add an option to toggle on/off


% set_conservations(?Ai, +Vars)
%
% given a list of variables for the species of the system, use the basis of
% P-invariants already computed to enforce constraints on the exponents Ai
set_conservations(Ai, Places) :-
  nb_getval(base, B),
  invariants:vars(Vars),
  length(Ai, N),
  length(Bi, N),
  reorder_place_vars(Ai, Bi, Places, Vars),
  get_initial_values(Vars, Di),
  set_conservations_rec(B, Bi, Di),
  set_constant_vars.


%% reorder_place_vars(-L1, -L2, +L3, +L4)
%
% reorders the list of vars L1, with order L3, as L2 with order L4
reorder_place_vars([], _, [], _).

reorder_place_vars([H | T], L, [HP | P], V) :-
  nth0(N, V, HP),
  nth0(N, L, H),
  !,
  reorder_place_vars(T, L, P, V).

%% set_constant_vars
%
% add fake parameters corresponding to initial values of constant variables
set_constant_vars :-
  invariants:vars(Vars),
  enumerate_molecules(M),
  member(V, M),
  get_initial_concentration(V, I),
  \+ member(V, Vars),
  get_degree(I, DI),
  assertz(degree(V, DI)),
  fail.

set_constant_vars.

% set_conservations_rec(-B, ?Ai, -Di)
%
% same as above but for a given basis B and already computed initial values Di
set_conservations_rec([], _, _).

set_conservations_rec([B | BB], Ai, Di) :-
  set_conservation(B, Ai, [], Di, 0),
  set_conservations_rec(BB, Ai, Di).


% set_conservation(-B, ?Ai, L, -Di, ?D)
%
% same as above but for a given vector B of the basis and using L as
% accumulator for storing the exponents that are in the support of B
% and D as accumulator for the initial value
set_conservation([0 | BB], [_ | AA], L, [_ | DD], D) :-
  !,
  set_conservation(BB, AA, L, DD, D).

set_conservation([N | BB], [Ai | AA], L, [Di | DD], D) :-
  N > 0,
  Dj is N * Di + D,
  set_conservation(BB, AA, [Ai | L], DD, Dj).

set_conservation([], [], L, [], D) :-
  conserve(L, D).


:- dynamic(degree/2).

% set_degrees(+Epsilon)
%
% given a small Epsilon, compute rescaling degrees for parameters
set_degrees(Epsilon) :-
  Epsilon < 1,
  nb_setval(epsilon, Epsilon),
  E is log(Epsilon),
  nb_setval(log_epsilon, E),
  retractall(degree(_, _)),
  parameter_value(K, V),
  get_degree(V, Deg),
  assertz(degree(K, Deg)),
  fail.

set_degrees(_).


get_degree(N, -inf) :-
  N =< 0,
  !.

get_degree(D, DD) :-
  nb_getval(log_epsilon, LEpsilon),
  nb_getval(ovidiu_denominator, O),
  % DD is round((1 + O) * log(D)/LEpsilon).
  DD is (1 + O) * round(log(D)/LEpsilon).

% subst_deg_and_vars(+E, +Vars, ?Ai, ?D)
%
% substitute parameters by their degree, species by their exponent and
% products by sums in E to get D
subst_deg_and_vars(A * B, Vars, Ai, AA + BB) :-
  !,
  subst_deg_and_vars(A, Vars, Ai, AA),
  subst_deg_and_vars(B, Vars, Ai, BB).

subst_deg_and_vars(A ^ N, Vars, Ai, N * AA) :-
  integer(N),
  !,
  subst_deg_and_vars(A, Vars, Ai, AA).

subst_deg_and_vars(A ^ N, Vars, Ai, NN * AA) :-
  float(N),
  NN is round(N),
  NN =:= N,
  !,
  subst_deg_and_vars(A, Vars, Ai, AA).

subst_deg_and_vars(N, _Vars, _Ai, D) :-
  number(N),
  !,
  get_degree(N, D).

subst_deg_and_vars(K, _Vars, _Ai, D) :-
  degree(K, D),
  !.

subst_deg_and_vars(V, Vars, Ai, A) :-
  nth0(N, Vars, V),
  !,
  nth0(N, Ai, A).

% Only accept division by constants
subst_deg_and_vars(A/N, Vars, Ai, AA - D) :-
  const_degree(N, D),
  !,
  subst_deg_and_vars(A, Vars, Ai, AA).

% Only accept division by constants
subst_deg_and_vars(N + M, _Vars, _Ai, D) :-
  const_degree(N+M, D),
  !.

subst_deg_and_vars(E, _, _, _) :-
  throw(error(unknown_kinetic_term(E))).


prolog:message(error(unknown_kinetic_term(E))) -->
  ['We do not understand the kinetic term ~p~n'-[E]].


const_degree(A, D) :-
  const_eval(A, V),
  get_degree(V, D).


const_eval(K, V) :-
  parameter_value(K, V),
  !.

const_eval(N, N) :-
  number(N),
  !.

const_eval(Term, D) :-
  Term =.. [Op, A, B],
  member(Op, ['+', '*', '-', '/']),
  const_eval(A, AA),
  const_eval(B, BB),
  TTerm =.. [Op, AA, BB],
  D is TTerm.

% split_pos_neg_and_subst(+O, +Vars, ?Ai, ?Pos, ?Neg)
%
% split expression O in positive and negative terms of the ODE
% for each term apply substitutions to keep the epsilon exponent
split_pos_neg_and_subst(O, Vars, Ai, Pos, Neg) :-
  split_pos_neg_and_subst_rec(O, Vars, Ai, Pos, Neg).


split_pos_neg_and_subst_rec(E + N * T, Vars, Ai, Pos, Neg) :-
  !,
  split_pos_neg_and_subst_rec(N * T, Vars, Ai, Pos1, Neg1),
  append(Pos1, PPos, Pos),
  append(Neg1, NNeg, Neg),
  split_pos_neg_and_subst_rec(E, Vars, Ai, PPos, NNeg).

split_pos_neg_and_subst_rec(N * (A - B), Vars, Ai, Pos, Neg) :-
  !,
  M is -N,
  split_pos_neg_and_subst_rec(N * A + M * B, Vars, Ai, Pos, Neg).

split_pos_neg_and_subst_rec(N * (A + B), Vars, Ai, Pos, Neg) :-
  !,
  split_pos_neg_and_subst_rec(N * A + N * B, Vars, Ai, Pos, Neg).

split_pos_neg_and_subst_rec(N * T, Vars, Ai, Pos, Neg) :-
  number(N),
  !,
  ode:substitute_functions(T, TT),
  subst_deg_and_vars(TT, Vars, Ai, S),
  (
    atom_is_subterm(inf, S)
  ->
    assertz(is_zero(T)),
    Pos = [],
    Neg = []
  ;
    (
      N > 0
    ->
      Neg = [],
      Pos = [S]
    ;
      Pos = [],
      Neg = [S]
    )
  ).

split_pos_neg_and_subst_rec(- E, Vars, Ai, Pos, Neg) :-
  !,
  split_pos_neg_and_subst_rec(E, Vars, Ai, Neg, Pos).

split_pos_neg_and_subst_rec(K, Vars, Ai, Pos, Neg) :-
  split_pos_neg_and_subst_rec(1*K, Vars, Ai, Pos, Neg).


% get_equilibrations(+Vars, ?Ai, ?E, ?DE)
%
% E is the list of equilibrations coming from the precomputed odes and taking
% into account that Ai are the exponents of species of Vars
% DE also substracts degree of left side for correct rescaled system
get_equilibrations(Vars, Ai, E, DE) :-
  ode_system,
  get_current_ode_system(Id),
  findall(
    (S, O),
    (
      item([parent: Id, kind: ode, item: I]),
      I = (d(S)/dt = O)
    ),
    L
  ),
  debug(tropical, "ODEs: ~w", [L]),
  get_equilibrations_rec(L, Vars, Ai, E, DE).


get_equilibrations_rec([], _, _, [], []).

get_equilibrations_rec([(S, O) | L], Vars, Ai, [E | EE], [DE | DEE]) :-
  nth0(N, Vars, S),
  nth0(N, Ai, A),
  split_pos_neg_and_subst(O, Vars, Ai, Pos, Neg),
  % since we will equalize mins, we do not really need to substract left and
  % right...
  substract_rec(Pos, A, PPos),
  substract_rec(Neg, A, NNeg),
  E = (Pos, Neg),
  DE = (PPos, NNeg),
  debug(tropical, "Monomials for ~w: ~w~n", [S, E]),
  get_equilibrations_rec(L, Vars, Ai, EE, DEE).


substract_rec(L, A, LA) :-
  maplist(tropical:reverse_minus(A), L, LA).


reverse_minus(A, B, B - A).


show_rescaling([], []) :-
  nl.

show_rescaling([A | AA], [V | VV]) :-
  epsilon_degree(-A, D),
  format("~w' = ~w~w~n", [V, D, V]),
  show_rescaling(AA, VV).


show_rescaling_ovidiu([D]) :-
  !,
  nb_getval(ovidiu_denominator, O),
  DD is D / (1 + O),
  write(ovidiu_file, DD),
  nl(ovidiu_file).

show_rescaling_ovidiu([H | T]) :-
  nb_getval(ovidiu_denominator, O),
  HH is H / (1 + O),
  write(ovidiu_file, HH),
  write(ovidiu_file, '\t'),
  show_rescaling_ovidiu(T).


epsilon_degree(D, A) :-
  (
    D =\= 0
  ->
    format(atom(A), "ε^(~p) * ", [D])
  ;
    A = ''
  ).


% look for an atom as a leaf of a term, not trying to unify FD vars
atom_is_subterm(_, Y) :-
  var(Y),
  !,
  fail.

atom_is_subterm(X, X) :-
  !.

atom_is_subterm(_, Y) :-
  atomic(Y),
  !,
  fail.

atom_is_subterm(X, F) :-
  F =.. [_ | Args],
  member(A, Args),
  atom_is_subterm(X, A).


%% molecules_vars(-L)
%
% returns variables of vars(L), i.e. found as constrained by invariants
% but ordered as in molecules(L).

molecules_vars(L) :-
  enumerate_molecules(P),
  invariants:vars(V),
  filter(P, V, L).


%% filter(+L1, +L2, +L3)
%
% L3 is elements of L1 that are also in L2
filter([], _, []).

filter([H | T], V, [H | L]) :-
  member(H, V),
  !,
  filter(T, V, L).

filter([_ | T], V, L) :-
  filter(T, V, L).


ignore_unbalanced([], []).

ignore_unbalanced([([_ | _], []) | L], LL) :-
  !,
  ignore_unbalanced(L, LL).


ignore_unbalanced([([], [_ | _]) | L], LL) :-
  !,
  ignore_unbalanced(L, LL).

ignore_unbalanced([H | T], [H | TT]) :-
  !,
  ignore_unbalanced(T, TT).


ignore_partial([], [], [], _).

ignore_partial([_ | E], Equils, [V | Vars], Ignored) :-
  memberchk(V, Ignored),
  !,
  ignore_partial(E, Equils, Vars, Ignored).


ignore_partial([Eq | E], [Eq | Equils], [_ | Vars], Ignored) :-
  ignore_partial(E, Equils, Vars, Ignored).


:- initial(option(tropical_epsilon: 0.1)).
:- initial(option(tropical_max_degree: 3)).
:- initial(option(tropical_ignore: {})).
:- initial(option(tropical_denominator: 0)).
:- initial(option(tropical_single_solution: no)).

:- biocham_silent(clear_model).


% main predicate for tropical equilibration
% if Max is a variable, use the number of variables as bound on the absolute
% value of the degree in epsilon, otherwise iterate powers of 2 up to 2^Max
tropicalize :-
  biocham_command,
  doc('Try to solve a tropical equilibration problem.
  \\begin{example}
  \\trace{
    biocham: load(library:examples/cell_cycle/Tyson_1991.bc).
    biocham: tropicalize(tropical_max_degree: 2).
  }
  \\end{example}'),
  option(
    tropical_epsilon,
    number,
    Eps,
    'Used for computing degrees (must be < 1)'
  ),
  option(
    tropical_max_degree,
    number,
    Max,
    'Max absolute value of the degree in epsilon as a power of 2'
  ),
  option(
    tropical_ignore,
    {object},
    Partial,
    'set of species to ignore for equilibration ({} to ignore none, {-1} to
     ignore all unbalanced)'
  ),
  option(
    tropical_denominator,
    number,
    Denominator,
    'Look for solutions that are integers/(tropical_denominator+1)'
  ),
  option(
    tropical_single_solution,
    yesno,
    Single,
    'Stop after finding one solution'
  ),
  option(stats, yesno, Stats, 'display computation time'),
  nb_setval(ovidiu_denominator, Denominator),
  set_degrees(Eps),
  retractall(is_zero(_)),
  (
    % might not be defined so we use nb_current instead of nb_getval
    nb_current(ovidiu_file, Ovidiu)
  ->
    open(Ovidiu, write, _, [alias(ovidiu_file)]),
    find_pinvar(5)
  ;
    find_pinvar(4)
  ),
  debug(tropical, "invar done", []),
  (
    Stats == yes
  ->
    statistics(cputime, Time0)
  ;
    true
  ),
  molecules_vars(Vars),
  length(Vars, VL),
  length(Ai, VL),
  (
    (
      (
        nb_current(ovidiu_file, _)
      ;
        Single = yes
      )
    ->
      MaxD is 2^Max
    ;
      between(2, Max, Power),
      MaxD is 2^Power
    ),
    debug(tropical, "about to set conservations… (Max = ~w)", [MaxD]),
    set_conservations(Ai, Vars),
    debug(tropical, "conservations done", []),
    get_equilibrations(Vars, Ai, E, _DE),
    debug(tropical, "equilibrations: ~w", [E]),
    (
      Partial == [{}]
    ->
      Equils = E
    ;
      Partial == [-1]
    ->
      ignore_unbalanced(E, Equils),
      write(E),nl,
      write(Equils),nl
    ;
      format("Ignoring equilibration for ~w\n", [Partial]),
      ignore_partial(E, Equils, Vars, Partial)
    ),
    debug(tropical, "equilibrations: ~w~n", [Equils]),
    equilibrate_rec(Equils),
    MinD is -MaxD,
    Ai ins MinD..MaxD,
    (
      Single = yes
    ->
      once(labeling([bisect], Ai))  % ffc makes biomodel 407 slow...
    ;
      labeling([bisect], Ai)        % bisect slightly better
    ),
    (
      nb_current(ovidiu_file, _)
    ->
      show_rescaling_ovidiu(Ai)
    ;
      write('\nfound a complete equilibration leading to the rescaling:\n'),
      show_rescaling(Ai, Vars)
    ),
    fail
  ;
    write('\nNo (more) complete equilibration\n'),
    (
      nb_current(ovidiu_file, _)
    ->
      close(ovidiu_file),
      nb_delete(ovidiu_file)
    ;
      Stats == yes
    ->
      statistics(cputime, Time1),
      Time is Time1 - Time0,
      format("Time: ~f~n~w variables~n", [Time, VL])
    ;
      true
    )
  ).