import sys
import platform
from decimal import Decimal

import numpy as np
from numpy.core import *
from numpy.random import rand, randint, randn
from numpy.testing import *
from numpy.testing.utils import WarningManager
from numpy.core.multiarray import dot as dot_
import warnings

class Vec:
    def __init__(self,sequence=None):
        if sequence is None:
            sequence=[]
        self.array=array(sequence)
    def __add__(self,other):
        out=Vec()
        out.array=self.array+other.array
        return out
    def __sub__(self,other):
        out=Vec()
        out.array=self.array-other.array
        return out
    def __mul__(self,other): # with scalar
        out=Vec(self.array.copy())
        out.array*=other
        return out
    def __rmul__(self,other):
        return self*other


class TestDot(TestCase):
    def setUp(self):
        self.A = rand(10,8)
        self.b1 = rand(8,1)
        self.b2 = rand(8)
        self.b3 = rand(1,8)
        self.b4 = rand(10)
        self.N = 14

    def test_matmat(self):
        A = self.A
        c1 = dot(A.transpose(), A)
        c2 = dot_(A.transpose(), A)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_matvec(self):
        A, b1 = self.A, self.b1
        c1 = dot(A, b1)
        c2 = dot_(A, b1)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_matvec2(self):
        A, b2 = self.A, self.b2
        c1 = dot(A, b2)
        c2 = dot_(A, b2)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecmat(self):
        A, b4 = self.A, self.b4
        c1 = dot(b4, A)
        c2 = dot_(b4, A)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecmat2(self):
        b3, A = self.b3, self.A
        c1 = dot(b3, A.transpose())
        c2 = dot_(b3, A.transpose())
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecmat3(self):
        A, b4 = self.A, self.b4
        c1 = dot(A.transpose(),b4)
        c2 = dot_(A.transpose(),b4)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecvecouter(self):
        b1, b3 = self.b1, self.b3
        c1 = dot(b1, b3)
        c2 = dot_(b1, b3)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecvecinner(self):
        b1, b3 = self.b1, self.b3
        c1 = dot(b3, b1)
        c2 = dot_(b3, b1)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_columnvect1(self):
        b1 = ones((3,1))
        b2 = [5.3]
        c1 = dot(b1,b2)
        c2 = dot_(b1,b2)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_columnvect2(self):
        b1 = ones((3,1)).transpose()
        b2 = [6.2]
        c1 = dot(b2,b1)
        c2 = dot_(b2,b1)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecscalar(self):
        b1 = rand(1,1)
        b2 = rand(1,8)
        c1 = dot(b1,b2)
        c2 = dot_(b1,b2)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecscalar2(self):
        b1 = rand(8,1)
        b2 = rand(1,1)
        c1 = dot(b1,b2)
        c2 = dot_(b1,b2)
        assert_almost_equal(c1, c2, decimal=self.N)

    def test_all(self):
        dims = [(),(1,),(1,1)]
        for dim1 in dims:
            for dim2 in dims:
                arg1 = rand(*dim1)
                arg2 = rand(*dim2)
                c1 = dot(arg1, arg2)
                c2 = dot_(arg1, arg2)
                assert_(c1.shape == c2.shape)
                assert_almost_equal(c1, c2, decimal=self.N)

    def test_vecobject(self):
        U_non_cont = transpose([[1.,1.],[1.,2.]])
        U_cont = ascontiguousarray(U_non_cont)
        x = array([Vec([1.,0.]),Vec([0.,1.])])
        zeros = array([Vec([0.,0.]),Vec([0.,0.])])
        zeros_test = dot(U_cont,x) - dot(U_non_cont,x)
        assert_equal(zeros[0].array, zeros_test[0].array)
        assert_equal(zeros[1].array, zeros_test[1].array)


class TestResize(TestCase):
    def test_copies(self):
        A = array([[1,2],[3,4]])
        Ar1 = array([[1,2,3,4],[1,2,3,4]])
        assert_equal(resize(A, (2,4)), Ar1)

        Ar2 = array([[1,2],[3,4],[1,2],[3,4]])
        assert_equal(resize(A, (4,2)), Ar2)

        Ar3 = array([[1,2,3],[4,1,2],[3,4,1],[2,3,4]])
        assert_equal(resize(A, (4,3)), Ar3)

    def test_zeroresize(self):
        A = array([[1,2],[3,4]])
        Ar = resize(A, (0,))
        assert_equal(Ar, array([]))

class TestNonarrayArgs(TestCase):
    # check that non-array arguments to functions wrap them in arrays
    def test_squeeze(self):
        A = [[[1,1,1],[2,2,2],[3,3,3]]]
        assert_(squeeze(A).shape == (3,3))

    def test_cumproduct(self):
        A = [[1,2,3],[4,5,6]]
        assert_(all(cumproduct(A) == array([1,2,6,24,120,720])))

    def test_size(self):
        A = [[1,2,3],[4,5,6]]
        assert_(size(A) == 6)
        assert_(size(A,0) == 2)
        assert_(size(A,1) == 3)

    def test_mean(self):
        A = [[1,2,3],[4,5,6]]
        assert_(mean(A) == 3.5)
        assert_(all(mean(A,0) == array([2.5,3.5,4.5])))
        assert_(all(mean(A,1) == array([2.,5.])))

    def test_std(self):
        A = [[1,2,3],[4,5,6]]
        assert_almost_equal(std(A), 1.707825127659933)
        assert_almost_equal(std(A,0), array([1.5, 1.5, 1.5]))
        assert_almost_equal(std(A,1), array([0.81649658, 0.81649658]))

    def test_var(self):
        A = [[1,2,3],[4,5,6]]
        assert_almost_equal(var(A), 2.9166666666666665)