#!/usr/bin/python3

# pip install ortools
# https://pypi.org/project/ortools/
from ortools.constraint_solver import pywrapcp

def main():
  solver = pywrapcp.Solver("PE_44")

  max_limit=10**7 # may be raised gradually

  P1=solver.IntVar(0, max_limit, "P1")
  P2=solver.IntVar(0, max_limit, "P2")
  P_diff=solver.IntVar(0, max_limit, "P_diff")
  P_sum=solver.IntVar(0, max_limit, "P_sum")

  n1=solver.IntVar(0, max_limit, "n1")
  n2=solver.IntVar(0, max_limit, "n2")
  n_diff=solver.IntVar(0, max_limit, "n_diff")
  n_sum=solver.IntVar(0, max_limit, "n_sum")

  solver.Add(P1>1)
  solver.Add(P2>1)
  solver.Add(P_diff>1)
  solver.Add(P_sum>1)

  solver.Add(P1!=P2)

  solver.Add(n1>1)
  solver.Add(n2>1)
  solver.Add(n_diff>1)
  solver.Add(n_sum>1)

  solver.Add(2*P1==n1*(3*n1-1))
  solver.Add(2*P2==n2*(3*n2-1))
  solver.Add(2*P_diff==n_diff*(3*n_diff-1))
  solver.Add(2*P_sum==n_sum*(3*n_sum-1))

  solver.Add(P1-P2==P_diff)
  solver.Add(P1+P2==P_sum)

  # objective
  objective = solver.Minimize(P_diff, 1)

  solution = solver.Assignment()

  db = solver.Phase([P1, P2, P_diff, P_sum, n1, n2, n_diff, n_sum],
                    solver.CHOOSE_MIN_SIZE_LOWEST_MAX,
                    solver.ASSIGN_MIN_VALUE)

  solver.NewSearch(db, [objective])
  assert solver.NextSolution()==True
  print(P1)
  print(P2)
  print(P_diff)
  print(P_sum)
  print(n1)
  print(n2)
  print(n_diff)
  print(n_sum)
  solver.EndSearch()

main()

