import unittest
from pylearn.algorithms.linear_regression import *
from pylearn.datasets.make_test_datasets import *
import numpy

class test_linear_regression(unittest.TestCase):

    def test1(self):
        trainset,testset,theta=make_artificial_datasets_from_function(n_inputs=3,
                                                                      n_targets=2,
                                                                      n_examples=100,
                                                                      f=linear_predictor)
        
        assert trainset.fields()['input'].shape==(50,3)
        assert testset.fields()['target'].shape==(50,2)
        regressor = LinearRegression(L2_regularizer=0.1)
        predictor = regressor(trainset)
        test_data = testset.fields()
        mse = predictor.compute_mse(test_data['input'],test_data['target'])
        print 'mse = ',mse
        
if __name__ == '__main__':
    import sys

    if len(sys.argv)==1:
        unittest.main()
    else:
        assert sys.argv[1]=="--debug"
        tests = []
        for arg in sys.argv[2:]:
            tests.append(arg)
        if tests:
            unittest.TestSuite(map(T_DataSet, tests)).debug()
        else:
            module = __import__("_test_linear_regression")
            tests = unittest.TestLoader().loadTestsFromModule(module)
            tests.debug()