#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license, see LICENSE.
"""
Tests for the skhep.math.vectors module.
"""

# -----------------------------------------------------------------------------
# Import statements
# -----------------------------------------------------------------------------
import pytest
from pytest import approx

from skhep.math.vectors import *
from skhep.utils.exceptions import *
from math import pi, sqrt


# -----------------------------------------------------------------------------
# Actual tests
# -----------------------------------------------------------------------------
def test_vectors_constructors():
    with pytest.raises(TypeError):
        Vector3D.fromiterable(1.0)
    with pytest.raises(ValueError):
        Vector3D.fromiterable([1.0])
    with pytest.raises(ValueError):
        Vector3D.fromiterable(["str1", "str2", "str3"])
    #
    v1 = Vector3D()
    assert str(v1) == str((0.0, 0.0, 0.0))
    assert repr(v1) == "<Vector3D (x=0.0,y=0.0,z=0.0)>"
    assert str(v1) == str(Vector3D.origin())
    v2 = Vector3D(1.0, 1.0, 1.0)
    assert str(v2) == str((1.0, 1.0, 1.0))
    v3 = Vector3D.fromvector(v1)
    assert str(v3) == str((0.0, 0.0, 0.0))
    v4 = Vector3D.fromiterable([1.0, 1.0, 1.0])
    assert str(v4) == str((1.0, 1.0, 1.0))
    v5 = Vector3D.fromcylindricalcoords(1.0, 0.0, 1.0)
    assert v5 == Vector3D(1.0, 0.0, 1.0)
    assert v5.rho == 1.0
    v6 = Vector3D.fromcylindricalcoords(0.5, pi / 2, 0.0)
    assert v6 == Vector3D(0.0, 0.5, 0.0)
    v7 = Vector3D.fromsphericalcoords(1.0, 0.0, 0.0)
    assert v7 == Vector3D(0.0, 0.0, 1.0)
    assert v7.r == 1.0
    assert v7.theta() == 0.0
    v8 = Vector3D.fromsphericalcoords(1.0, 0.0, pi / 2)
    assert v8 == Vector3D(0.0, 0.0, 1.0)
    v9 = Vector3D.fromsphericalcoords(2.0, pi / 2, pi / 4)
    assert v9 == Vector3D(sqrt(2), sqrt(2), 0.0)
    v10 = Vector3D.frompoint(2.0, 2.0, 2.0)
    assert str(v10) == str((2.0, 2.0, 2.0))
    #
    with pytest.raises(TypeError):
        LorentzVector.fromiterable(1)
    with pytest.raises(ValueError):
        LorentzVector.fromiterable([1])
    with pytest.raises(ValueError):
        LorentzVector.fromiterable(["str1", "str2", "str3", "str4"])
    #
    lv1 = LorentzVector()
    assert str(lv1) == str((0.0, 0.0, 0.0, 0.0))
    assert repr(lv1) == "<LorentzVector (x=0.0,y=0.0,z=0.0,t=0.0)>"
    lv2 = LorentzVector(1.0, 1.0, 1.0, 1.0)
    assert str(lv2) == str((1.0, 1.0, 1.0, 1.0))
    lv3 = LorentzVector.from4vector(lv1)
    assert str(lv3) == str((0.0, 0.0, 0.0, 0.0))
    lv4 = LorentzVector.from3vector(v2, 2.0)
    assert str(lv4) == str((1.0, 1.0, 1.0, 2.0))
    lv5 = LorentzVector.fromiterable([1.0, 1.0, 1.0, 0.0])
    assert str(lv5) == str((1.0, 1.0, 1.0, 0.0))


def test_containers_properties():
    with pytest.raises(IndexError):
        Vector3D.__setitem__(Vector3D(), 3, 1.0)
    with pytest.raises(IndexError):
        Vector3D.__getitem__(Vector3D(), 3)
    #
    v1 = Vector3D()
    v1.x = 5.0
    v1.y = -5.0
    v1.z = 10
    assert v1 == Vector3D(5.0, -5.0, 10.0)
    v1.set(-5.0, 5.0, 10.0)
    assert v1 == Vector3D(-5.0, 5.0, 10.0)
    v1[0] = 1.0
    assert v1 == Vector3D(1.0, 5.0, 10.0)
    v1[1] = 1.0
    assert v1 == Vector3D(1.0, 1.0, 10.0)
    v1[2] = 1.0
    assert v1 == Vector3D(1.0, 1.0, 1.0)
    assert v1[0] == 1.0
    assert v1[1] == 1.0
    assert v1[2] == 1.0
    assert len(v1) == 3.0
    assert v1.tolist() == [1.0, 1.0, 1.0]
    assert list(v1) == [1.0, 1.0, 1.0]
    assert [v for v in v1] == [1.0, 1.0, 1.0]
    #
    with pytest.raises(IndexError):
        LorentzVector.__setitem__(LorentzVector(), 4, 1.0)
    with pytest.raises(IndexError):
        LorentzVector.__getitem__(LorentzVector(), 4)
    #
    lv1 = LorentzVector()
    lv1.x = 5.0
    lv1.y = -5.0
    lv1.z = 10
    lv1.t = 2.0
    assert lv1 == LorentzVector(5.0, -5.0, 10.0, 2.0)
    lv1.px = 5.0
    lv1.py = 5.0
    lv1.pz = -10
    lv1.e = 2.0
    assert lv1 == LorentzVector(5.0, 5.0, -10.0, 2.0)
    lv1.set(-5.0, 5.0, 10.0, -2.0)
    assert lv1 == LorentzVector(-5.0, 5.0, 10.0, -2.0)
    lv1[0] = 1.0
    assert lv1 == LorentzVector(1.0, 5.0, 10.0, -2.0)
    lv1[1] = 1.0
    assert lv1 == LorentzVector(1.0, 1.0, 10.0, -2.0)
    lv1[2] = 1.0
    assert lv1 == LorentzVector(1.0, 1.0, 1.0, -2.0)
    lv1[3] = 1.0
    assert lv1 == LorentzVector(1.0, 1.0, 1.0, 1.0)
    assert lv1[0] == 1.0
    assert lv1[1] == 1.0
    assert lv1[2] == 1.0
    assert lv1[3] == 1.0
    assert len(lv1) == 4.0
    assert lv1.tolist() == [1.0, 1.0, 1.0, 1.0]
    assert list(lv1) == [1.0, 1.0, 1.0, 1.0]
    assert [v for v in lv1] == [1.0, 1.0, 1.0, 1.0]


def test_vectors_operators():
    with pytest.raises(InvalidOperationError):
        Vector3D.__iadd__(Vector3D(), "a")
    with pytest.raises(InvalidOperationError):
        Vector3D.__isub__(Vector3D(), "a")
    with pytest.raises(InvalidOperationError):
        Vector3D.__imul__(Vector3D(), Vector3D())
    with pytest.raises(ZeroDivisionError):
        Vector3D.__div__(Vector3D(), 0.0)
    with pytest.raises(InvalidOperationError):
        Vector3D.__idiv__(Vector3D(), Vector3D())
    with pytest.raises(ZeroDivisionError):
        Vector3D.__idiv__(Vector3D(), 0.0)
    #
    v1, v2 = Vector3D(0.0, 0.0, 0.0), Vector3D(1.0, 1.0, 1.0)
    v3, v4 = Vector3D(2.0, 2.0, 2.0), Vector3D(3.0, 3.0, 3.0)
    v5 = Vector3D(1.0, 2.0, 3.0)
    assert v1 == 0.0
    assert v1 + v2 == Vector3D(1.0, 1.0, 1.0)
    assert v1 - v2 == Vector3D(-1.0, -1.0, -1.0)
    assert v1 - v2 == -1.0 * Vector3D(1.0, 1.0, 1.0)
    assert v1 * 2.0 == Vector3D(0.0, 0.0, 0.0)
    assert v2 * 2.0 == Vector3D(2.0, 2.0, 2.0)
    assert 2.0 * v1 == Vector3D(0.0, 0.0, 0.0)
    assert 2.0 * v2 == Vector3D(2.0, 2.0, 2.0)
    assert v1 * v2 == 0.0
    assert v2 * v1 == 0.0
    assert v3 * v4 == v3.dot(v4), "Check commutativity"
    assert v3 * v4 == v4.dot(v3), "Check commutativity"
    assert v4 * v3 == v3.dot(v4), "Check commutativity"
    assert v4 * v3 == v4.dot(v3), "Check commutativity"
    assert v3 * v4 == 18.0
    assert v3 / 2.0 == v2
    assert v4 / 3.0 == v2
    v1 *= 2.0
    v2 *= 2.0
    assert v1 == Vector3D(0.0, 0.0, 0.0)
    assert v2 == Vector3D(2.0, 2.0, 2.0)
    v1 /= 2.0
    v2 /= 2.0
    assert v1 == Vector3D(0.0, 0.0, 0.0)
    assert v2 == Vector3D(1.0, 1.0, 1.0)
    assert (v1 + v2) * v3 == 6.0, "Check operations combination"
    assert (v2 - v1) * v3 == 6.0, "Check operations combination"
    assert v3 * (v1 - v2) == -6.0, "Check operations combination"
    assert 18.0 / (v3 * v4) == 1.0, "Check operations combination"
    assert v4 / (v3 * v2) == Vector3D(0.5, 0.5, 0.5)
    assert v2.cross(v2) == Vector3D(0.0, 0.0, 0.0)
    assert v2.cross(v5) == Vector3D(1.0, -2.0, 1.0)
    assert v5.cross(v2) == -1 * v2.cross(v5)
    assert v2 == v2
    assert v2 != v1
    assert not v2 != v2
    assert not v2 == v1
    assert not v1 == "a"
    assert not v1 == 1
    #
    with pytest.raises(InvalidOperationError):
        LorentzVector.__iadd__(LorentzVector(), "a")
    with pytest.raises(InvalidOperationError):
        LorentzVector.__isub__(LorentzVector(), "a")
    with pytest.raises(InvalidOperationError):
        LorentzVector.__imul__(LorentzVector(), LorentzVector())
    with pytest.raises(ZeroDivisionError):
        LorentzVector.__div__(LorentzVector(), 0.0)
    with pytest.raises(InvalidOperationError):
        LorentzVector.__idiv__(LorentzVector(), LorentzVector())
    with pytest.raises(ZeroDivisionError):
        LorentzVector.__idiv__(LorentzVector(), 0.0)
    #
    lv1, lv2 = LorentzVector(0.0, 0.0, 0.0, 0.0), LorentzVector(1.0, 1.0, 1.0, 0.0)
    lv3, lv4 = LorentzVector(2.0, 2.0, 2.0, 1.0), LorentzVector(3.0, 3.0, 3.0, 1.0)
    lv5 = LorentzVector(1.0, 1.0, 1.0, 6.0)
    assert lv1 == 0.0
    assert lv1 + lv2 == LorentzVector(1.0, 1.0, 1.0, 0.0)
    assert lv1 - lv2 == LorentzVector(-1.0, -1.0, -1.0, 0.0)
    assert lv1 * 2.0 == LorentzVector(0.0, 0.0, 0.0, 0.0)
    assert lv2 * 2.0 == LorentzVector(2.0, 2.0, 2.0, 0.0)
    assert 2.0 * lv1 == LorentzVector(0.0, 0.0, 0.0, 0.0)
    assert 2.0 * lv2 == LorentzVector(2.0, 2.0, 2.0, 0.0)
    assert lv2 * lv1 == 0.0
    assert lv1 * lv2 == 0.0
    assert lv3 * lv4 == lv3.dot(lv4)
    assert lv3 * lv4 == lv4.dot(lv3)
    assert lv4 * lv3 == lv3.dot(lv4)
    assert lv4 * lv3 == lv4.dot(lv3)
    assert lv3 * lv4 == -17.0
    assert lv3 / 2.0 == LorentzVector(1.0, 1.0, 1.0, 0.5)
    assert lv4 / 3.0 == LorentzVector(1.0, 1.0, 1.0, 1.0 / 3)
    assert lv3 * lv5 == 0.0
    lv1 *= 2.0
    lv2 *= 2.0
    assert lv1 == LorentzVector(0.0, 0.0, 0.0, 0.0)
    assert lv2 == LorentzVector(2.0, 2.0, 2.0, 0.0)
    lv1 /= 2.0
    lv2 /= 2.0
    assert lv1 == LorentzVector(0.0, 0.0, 0.0, 0.0)
    assert lv2 == LorentzVector(1.0, 1.0, 1.0, 0.0)
    assert lv2 == lv2
    assert lv2 != lv1
    assert not lv2 != lv2
    assert not lv2 == lv1
    assert not lv1 == "a"
    assert not lv1 == 1


def test_vectors_rotations():
    with pytest.raises(TypeError):
        Vector3D.rotate(Vector3D(), pi, 1)
    with pytest.raises(TypeError):
        Vector3D.rotate(Vector3D(), pi, [1, 2])
    with pytest.raises(TypeError):
        Vector3D.rotate(Vector3D(), pi, 1, 2, 3, 4)
    with pytest.raises(ValueError):
        Vector3D.rotate(Vector3D(), pi, 0, 1, "a")
    with pytest.raises(ValueError):
        Vector3D.rotate(Vector3D(), pi, ["a", "b", 3])
    #
    v1 = Vector3D.fromcylindricalcoords(1.0, 0.0, 0.0)
    assert v1.phi() == 0.0
    assert v1.rotatez(pi / 2).phi() == pi / 2
    assert v1.rotatez(pi / 2) == Vector3D(0.0, 1.0, 0.0)
    assert v1.rotatey(pi).phi() == pi
    assert v1.rotatey(-pi).phi() == pi
    assert v1.rotatey(pi) == Vector3D(-1.0, 0.0, 0.0)
    assert v1.rotatey(-pi) == Vector3D(-1.0, 0.0, 0.0)
    assert v1.rotatex(pi).phi() == 0.0
    assert v1.rotatex(pi) == Vector3D(1.0, 0.0, 0.0)
    v2 = Vector3D.fromsphericalcoords(1.0, pi / 2, pi / 2)
    assert v2.phi() == pi / 2
    assert v2.theta() == pi / 2
    assert v2.rotatex(pi).phi() == -pi / 2
    assert v2.rotatex(pi).theta() == pi / 2
    assert v2.rotatex(pi) == Vector3D(0.0, -1.0, 0.0)
    v3 = Vector3D.fromsphericalcoords(1.0, pi / 4, pi / 4)
    angle = v2.angle(v3)
    axis = v2.cross(v3)
    assert v2.rotate(angle, axis) == v3
    assert v2.rotate(-angle, -1.0 * axis) == v3
    #
    with pytest.raises(TypeError):
        LorentzVector.rotate(LorentzVector(), pi, 1)
    with pytest.raises(TypeError):
        LorentzVector.rotate(LorentzVector(), pi, [1, 2])
    with pytest.raises(TypeError):
        LorentzVector.rotate(LorentzVector(), pi, 1, 2, 3, 4)
    with pytest.raises(ValueError):
        LorentzVector.rotate(LorentzVector(), pi, 0, 1, "a")
    with pytest.raises(ValueError):
        LorentzVector.rotate(LorentzVector(), pi, ["a", "b", 3])
    #
    lv1 = LorentzVector.from3vector(v1, 1.0)
    assert lv1.phi() == 0.0
    assert lv1.rotatez(pi / 2).phi() == pi / 2
    assert lv1.rotatez(pi / 2) == LorentzVector(0.0, 1.0, 0.0, 1.0)
    assert lv1.rotatey(pi).phi() == pi
    assert lv1.rotatey(-pi).phi() == pi
    assert lv1.rotatey(pi) == LorentzVector(-1.0, 0.0, 0.0, 1.0)
    assert lv1.rotatey(-pi) == LorentzVector(-1.0, 0.0, 0.0, 1.0)
    assert lv1.rotatex(pi).phi() == 0.0
    assert lv1.rotatex(pi) == LorentzVector(1.0, 0.0, 0.0, 1.0)
    lv2 = LorentzVector.from3vector(v2, 2.0)
    assert lv2.phi() == pi / 2
    assert lv2.theta() == pi / 2
    assert lv2.rotatex(pi).phi() == -pi / 2
    assert lv2.rotatex(pi).theta() == pi / 2
    assert lv2.rotatex(pi) == LorentzVector(0.0, -1.0, 0.0, 2.0)
    lv3 = LorentzVector.from3vector(v3, 2.0)
    assert lv2.rotate(angle, axis) == lv3
    assert lv2.rotate(-angle, -1.0 * axis) == lv3


def test_3Dvectors_properties():
    v0 = Vector3D()
    v1, v2 = Vector3D(1.0, 1.0, 1.0), Vector3D(2.0, 2.0, 2.0)
    v3, v4 = Vector3D(-1.0, -1.0, -1.0), Vector3D(-2.0, -2.0, -2.0)
    v5, v6 = Vector3D(1.0, 1.0, 0.0), Vector3D(0.0, 0.0, 2.0)
    assert not v0
    assert v1
    assert v1.mag2 == 3.0
    assert v1.mag == sqrt(3.0)
    assert v2.mag2 == 12.0
    assert v2.mag == sqrt(12.0)
    assert v1.unit().mag == 1.0
    v7 = v1.unit()
    assert v1.unit() == v7.unit()
    assert v1.isparallel(v2) == True
    assert v2.isparallel(v1) == True
    assert v1.isantiparallel(v2) == False
    assert v2.isantiparallel(v1) == False
    assert v1.isantiparallel(v3) == True
    assert v2.isantiparallel(v4) == True
    assert v1.isparallel(v3) == False
    assert v2.isparallel(v4) == False
    assert v1.isopposite(v3) == True
    assert v1.isopposite(v4) == False
    assert v2.isopposite(v4) == True
    assert v2.isopposite(v3) == False
    assert v1.isperpendicular(v2) == False
    assert v2.isperpendicular(v1) == False
    assert v5.isperpendicular(v6) == True
    assert v6.isperpendicular(v5) == True


def test_lorentz_vectors_properties():
    with pytest.raises(TypeError):
        LorentzVector.boost(LorentzVector(), 1)
    with pytest.raises(TypeError):
        LorentzVector.boost(LorentzVector(), [1, 2])
    with pytest.raises(TypeError):
        LorentzVector.boost(LorentzVector(), 1, 2, 3, 4)
    with pytest.raises(ValueError):
        LorentzVector.boost(LorentzVector(), 0, 1, "a")
    with pytest.raises(ValueError):
        LorentzVector.boost(LorentzVector(), ["a", "b", 3])
    #
    lv0 = LorentzVector()
    lv1, lv2 = LorentzVector(1.0, 1.0, 1.0, 1.0), LorentzVector(1.0, 1.0, 1.0, 2.0)
    assert lv1.boostvector == Vector3D(1.0, 1.0, 1.0)
    assert lv1.beta == sqrt(3.0)
    assert lv2.boostvector == Vector3D(0.5, 0.5, 0.5)
    lv3 = LorentzVector(0.0, 0.0, 1.0, 0.0)
    beta = 0.05
    gamma = 1 / sqrt(1 - beta**2)
    lv4 = lv3.boost(0, 0, beta)
    assert lv4.z == lv3.z * gamma, "Check length contraction"
    assert lv4.x == lv3.x
    assert lv4.y == lv3.x
    assert lv4.t == gamma * (lv3.t - beta * lv3.z)
    assert lv4 == LorentzVector(
        lv3.x, lv3.y, lv3.z * gamma, gamma * (lv3.t - beta * lv3.z)
    )
    lv5 = LorentzVector(0.0, 0.0, 0.0, 1.0)
    assert lv5.beta == 0.0
    assert lv5.gamma == 1.0
    assert lv5.boost(0, 0, 0) == lv5
    lv6 = lv5.boost([0, beta, 0])
    assert lv6.x == lv5.x
    assert lv6.z == lv5.z
    assert lv6.t == lv5.t * gamma, "Check time dilation"
    assert lv6.y == gamma * (lv5.y - beta * lv5.t)
    assert lv6 == LorentzVector(
        lv5.x, gamma * (lv5.y - beta * lv5.t), lv5.z, lv5.t * gamma
    )
    assert lv1.isspacelike() == True
    assert lv1.istimelike() == False
    assert lv1.islightlike() == False
    assert lv2.isspacelike() == False
    assert lv2.istimelike() == True
    assert lv2.islightlike() == False
    lv7 = LorentzVector(1.0, 1.0, 1.0, sqrt(3.0))
    assert lv7.isspacelike() == False
    assert lv7.istimelike() == False
    assert lv7.islightlike() == True
    #
    p1 = LorentzVector()
    p1.setpxpypzm(5.0, 5.0, 10.0, 5)
    assert p1.px == 5.0
    assert p1.py == 5.0
    assert p1.pz == 10.0
    assert p1.m == 5.0
    assert p1.mass == 5.0
    assert p1.m2 == 25.0
    assert p1.mass2 == 25.0
    assert p1.pt == sqrt(p1.px**2 + p1.py**2)
    assert p1.p == sqrt(p1.px**2 + p1.py**2 + p1.pz**2)
    assert p1.p == sqrt(p1.pt**2 + p1.pz**2)
    assert p1.e == sqrt(p1.m**2 + p1.p**2)
    assert p1.beta == p1.p / p1.e
    p2 = p1.boost(p1.px / p1.e, p1.y / p1.e, p1.z / p1.e)
    assert p2 == p1.boost(p1.boostvector)
    assert p2.p == 0.0
    assert p2.m == approx(5.0)
    p3 = LorentzVector()
    p3.setpxpypze(5.0, 5.0, 10.0, 20)
    assert p3.px == 5.0
    assert p3.py == 5.0
    assert p3.pz == 10.0
    assert p3.e == 20.0
    assert p3.m == sqrt(p3.e**2 - p3.p**2)
    assert p3.beta == p3.p / p3.e
    p4 = LorentzVector()
    p4.setptetaphim(10.0, 3.5, pi / 3, 5.0)
    assert p4.pt == approx(10.0)
    assert p4.eta == approx(3.5)
    assert p4.phi() == approx(pi / 3)
    assert p4.m == approx(5.0)
    p5 = LorentzVector()
    p5.setptetaphie(10.0, 3.9, -2 * (pi / 3), 20.0)
    assert p5.pt == approx(10.0)
    assert p5.eta == approx(3.9)
    assert p5.pseudorapidity == approx(3.9)
    assert p5.phi() == approx(-2 * (pi / 3))
    assert p5.e == approx(20.0)
    assert p5.deltaeta(p4) == approx(0.4)
    assert p4.deltaeta(p5) == approx(-0.4)
    assert p5.deltaphi(p4) == approx(-pi)
    assert p4.deltaphi(p5) == approx(pi)
    assert p5.deltar(p4) == approx(sqrt(0.4**2 + pi**2))
    assert p4.deltar(p5) == approx(sqrt(0.4**2 + pi**2))
    p6 = LorentzVector()
    p6.setptetaphie(10.0, 3.9, -pi, 20.0)
    assert p6.deltaphi(p4) == approx(2 * (pi / 3))
    assert p4.deltaphi(p6) == approx(-2 * (pi / 3))
    p7 = LorentzVector()
    p7.setpxpypzm(5.0, 5.0, 5.0, 0.0)
    assert p7.beta == 1.0
    assert p7.gamma == 10e10, "Gamma of the photons is +inf"
    assert p7.p == p7.e, "Momentum = Energy for photons"
    assert p7.pseudorapidity == p7.rapidity
    assert p7.pt == p7.et, "Transverse Momentum = Transverse Energy for photons"
    assert p7.pt == approx(p7.mt), "Transverse Momentum = Transverse Mass for photons"
    assert p7.perp2 == approx(p7.mt2)
    assert p7.islightlike() == True
    p8 = p1 + p7
    p9, p10 = p1.boost(p8.boostvector), p7.boost(p8.boostvector)
    assert p9.p == approx(p10.p), "Check boost to the C.O.M frame"
    p3_9, p3_10 = p9.vector, p10.vector
    assert p3_9.isopposite(p3_10)
    p8 = LorentzVector()
    p8.setptetaphie(10.0, 2e2, -2 * (pi / 3), 20.0)  # sinh(eta) diverge quickly
    assert p8.theta() == 0.0
    assert p8.eta == 10e10
