#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright © 2010-11 Mayo Clinic College of Medicine
#
# License: GNU General Public License, Version 3
# <http://www.gnu.org/licenses/gpl.html>
#
# Author: Donnie Berkholz <dberkholz@gentoo.org>

import scipy.odr
import scipy.stats
import numpy
import sys

if len(sys.argv) > 1:
    infiles = sys.argv[1:]
else:
    print '''
odr_fit.py FILE [FILE...]

Performs an orthogonal regression (a.k.a. ordinary least squares,
errors-in-variables) on two-column input files. This type of regression allows
for errors in both the X and Y variables, whereas typical linear regression
only allows for errors in Y. The results also show comparisons with linear
regressions using either X or Y as the dependent variable.
'''
    sys.exit(1)

initial_m = 1
initial_b = 0
initial_beta = [initial_m, initial_b]

def fcn(beta, x):
    m, b = beta

    return m*x + b

def read_data(infile):
    x = []
    y = []
    for line in open(infile):
        words = line.split()
        x.append(float(words[0]))
        y.append(float(words[1]))
    return x, y
                   
def linear_fit(x, y):
    final_m, final_b, r, prob, stderr = scipy.stats.stats.linregress(x, y)
    #print 'Linear\t\t%.2f\t%.2f\t%.2f' % (final_m, final_b, r**2),
    return final_m, final_b, r**2

def odr_fit(x, y):
    data = scipy.odr.Data(x, y)
    model = scipy.odr.Model(fcn=fcn)
    odr_obj = scipy.odr.ODR(data=data, model=model, beta0=initial_beta)
    output = odr_obj.run()

    final_m, final_b = output.beta
    r2 = 1.0 - output.res_var/numpy.var(y)

    #print 'Orthogonal\t%.2f\t%.2f\t%.2f' % (final_m, final_b, r2),
    return final_m, final_b, r2

def main():
    print 'Regression\tm\tb\tr^2\tFile'
    print
    for infile in infiles:
        x, y = read_data(infile)

        m, b, r2 = linear_fit(x, y)
        print 'Linear\t\t%.2f\t%.2f\t%.2f\t%s' % (m, b, r2, infile)

        m, b, r2 = linear_fit(y, x)
        print 'Inverse Linear\t%.2f\t%.2f\t%.2f\t%s' % (1/m, -b/m, r2, infile)

        m, b, r2 = odr_fit(x, y)
        print 'Orthogonal\t%.2f\t%.2f\t%.2f\t%s' % (m, b, r2, infile)

        print

main()