#!/usr/bin/env python2.2
########################################
# The contents of this file are subject to the MLX PUBLIC LICENSE version
# 1.0 (the "License"); you may not use this file except in
# compliance with the License.
# 
# Software distributed under the License is distributed on an "AS IS"
# basis, WITHOUT WARRANTY OF ANY KIND, either express or implied.  See
# the License for the specific language governing rights and limitations
# under the License.
# 
# The Original Source Code is "compClust", released 2003 September 03.
# 
# The Original Source Code was developed by the California Institute of
# Technology (Caltech).  Portions created by Caltech are Copyright (C)
# 2002-2003 California Institute of Technology. All Rights Reserved.
########################################
import unittest
import os

import Numeric
import RandomArray
import compClust.mlx.DA as DA
try:
  import compClust.mlx.cDA as cDA
  cDAexists = 1
except ImportError:
  cDAexists = 0

from compClust.mlx.datasets import Dataset
from compClust.mlx.labelings import Labeling
from compClust.mlx import models

import compClust

#import TestConstants

class DAmoduleTestCases(unittest.TestCase):
#class DATestCases:
  def setUp(self):
    self.original_dir = os.getcwd()
    os.chdir(compClust.mlx.__path__[0])

  def tearDown(self):
    os.chdir(self.original_dir)

  def setUpMatrix(self, num_vectors, num_classes=1):
    """Construct a model to play with"""
    # FIXME: Where should creating a simulated model live
    # FIXME: in setUp where it gets called for each test
    # FIXME: or in init where it's constructed just once.
    cluster_size =  num_vectors/num_classes
    self.means = []
    self.covars = []
    self.means.append(Numeric.array([0, 0, 0]))
    self.covars.append(Numeric.array([[1, .5, .1],
                                      [.5, 1, .5],
                                      [.1, .5, 1],]))
    self.means.append(Numeric.array([-1,-1,-1]))
    self.covars.append( Numeric.array([[.5, .1, 0],
                                       [.1, .5, .1],
                                       [0, .1, .5]]))
    
    self.means.append(Numeric.array([1, 2, 3]))
    self.covars.append(Numeric.array([[1, .1, .1],
                                      [.1, 1, .1],
                                      [ .1, .1, 1],]))
    
    self.means = Numeric.asarray(self.means)
    self.covars = Numeric.asarray(self.covars)
    
    self.data = Numeric.zeros((cluster_size*num_classes,self.covars[0].shape[0]),'d')
    self.classes = Numeric.zeros((cluster_size*num_classes),'i')
    
    for cluster in xrange(num_classes):
      for point in xrange(cluster_size):
        index = cluster * cluster_size + point
        self.data[index] = RandomArray.multivariate_normal(self.means[cluster],
                                                           self.covars[cluster])
        self.classes[index] = cluster

    self.dataset = Dataset(self.data)
    self.labels = Labeling(self.dataset)
    self.labels.labelRows(self.classes)

  def test_covar_weight_estimate(self):
    """Verify that covar_weight works correctly."""
    
    self.setUpMatrix(7500)
    data = self.dataset.getData()

    means = models.compute_model_means(self.dataset, self.labels)

    results = DA.covar_weights_estimate(data, means, self.classes)

    covars, weights = results

    # Compute measured - expected/expected for the vector
    variance = (covars - self.covars[0])/self.covars[0]

    for i in range(len(variance)):
      for j in range(len(variance[0])):
        # FIXME: is .1 a reasonable value to be using as a bound?
        assert variance[i][j] < .1, \
        "There was too much variance between the estimated covariance" + \
        "matrix and the origial matrix."

    if cDAexists:
      cresults = cDA.covar_weights_estimate(data, means, self.classes)
      self.failUnless(Numeric.allclose(covars, cresults[0]))
      self.failUnless(Numeric.allclose(weights, cresults[1]))
        

  def test_covariance_estimate(self):
    """Verify the covariance estimator works correctly.
    """
    if not cDAexists:
      return 
    
    self.setUpMatrix(750)
    data = self.dataset.getData()
    
    classes = map(self.labels.getLabelsByRow, range(self.dataset.getNumRows()))
    classes = map(lambda x : int(x[0]), classes)
    
    means = models.compute_model_means(self.dataset, self.labels)
    weights = models.compute_model_weights(self.dataset, self.labels)

    covars = DA.covariance_estimate(data, means, classes)

    # Compute measured - expected/expected for the vector
    variance = (covars - self.covars[0])/self.covars[0]

    for i in range(len(variance)):
      for j in range(len(variance[0])):
        # FIXME: is .1 a reasonable value to be using as a bound?
        assert variance[i][j] < .1, \
        "There was too much variance between the estimated covariance" + \
        "matrix and the origial matrix."

  def test_isposdef(self):
    pass

  def test_fix_covariance_matrix(self):
    pass

  def test_make_diag_matrix_square(self):
    m_start = [[ 1.0,  2.0,  3.0,  4.0,],
               [ 5.0,  6.0,  7.0,  8.0,],
               [ 9.0, 10.0, 11.0, 12.0,],
               [13.0, 14.0, 15.0, 16.0,],]
    m_start = Numeric.array(m_start)
    m_done  = [[ 1.0,  0.0,  0.0,  0.0,],
               [ 0.0,  6.0,  0.0,  0.0,],
               [ 0.0,  0.0, 11.0,  0.0,],
               [ 0.0,  0.0,  0.0, 16.0,],]
    m_done  = Numeric.array(m_done)
    
    new_matrix = DA.make_diag_matrix(m_start)

    if new_matrix != m_done:
      self.fail("Make diag matrix couldn't delete the off-diagonal elements in a 4x4 matrix")

  def test_vector_prob_gauss(self):
    self.setUpMatrix(10,1)

    daProbs = DA.vector_prob_gauss(self.data, self.means[0], self.covars[0])
    if cDAexists:
      cdaProbs = cDA.vector_prob_gauss(self.data, self.means[0], self.covars[0])
      self.failUnless(Numeric.allclose(daProbs, cdaProbs))
    
    
  def test_mixture_likelihood(self):
    self.setUpMatrix(3000,3)
    daCovars, daWeights = DA.covar_weights_estimate(self.data, self.means, self.classes)
    daLikelihood = DA.mixture_likelihood(self.data, self.means, daCovars, daWeights)

    # .1 is pretty pathetic, but it takes a lot of points to tighten
    # the estimate
    #self.failUnless(Numeric.allclose(daCovars, covars, rtol=.1))

    if cDAexists:
      cdaCovars, cdaWeights = cDA.covar_weights_estimate(self.data, self.means, self.classes)
      self.failUnless(Numeric.allclose(cdaCovars, daCovars))
      self.failUnless(Numeric.allclose(cdaWeights, daWeights))
      
      cdaLikelihood = cDA.mixture_likelihood(self.data, self.means, cdaCovars, cdaWeights)
      self.failUnless(Numeric.allclose(daLikelihood, cdaLikelihood))

  def test_isposdef_false(self):
    matrix = Numeric.array(
       [[  5.13388328e-05, -3.45171133e-05, -3.93597171e-05,],
        [ -3.45171133e-05,  2.32072106e-05,  2.64630833e-05,],
        [ -3.93597171e-05,  2.64630833e-05,  3.01757411e-05,],])

    if cDAexists:
      if DA.isposdef(matrix):
        self.fail("isposdef missed a positive indefinate matrix")

  def test_isposdef_true(self):
    matrix = Numeric.array(
       [[ 0.00343108, 0.00047888, 0.00045181,],
        [ 0.00047888, 0.00129044,-0.00012327,],
        [ 0.00045181,-0.00012327, 0.00063942,],])
    if cDAexists:
      if not DA.isposdef(matrix):
        self.fail("isposdef failed a positive definiate matrix")
  
def suite(**kw):
  RandomArray.seed( 0xf00f, 0xde1 )
  suite = unittest.makeSuite(DAmoduleTestCases)
  return suite

if __name__ == "__main__":
  unittest.main(defaultTest="suite")