#!/usr/bin/env python

import os
import matplotlib.pyplot as plt
import numpy as np
import kcorrect

rd = kcorrect.response.ResponseDict()

filename = os.path.join(kcorrect.KCORRECT_DIR,
                        'data', 'templates',
                        'kcorrect-lrg1-v4.fits')
templates = kcorrect.template.SED(filename=filename)

responses = ['sdss_u0', 'sdss_g0', 'sdss_r0', 'sdss_i0', 'sdss_z0',
             'twomass_Ks', 'wise_w4']

fitter = kcorrect.fitter.Fitter(templates=templates,
                                responses=responses, redshift_range=[0., 0.1],
                                nredshift=100)

fitter.set_Amatrix()

leff = np.array([rd[x].lambda_eff for x in responses], dtype=np.float32)

redshift = 0.0532
maggies = np.array([1., 1., 1., 1., 1., 1., 1.], dtype=np.float32)
ivar = np.array([1., 1., 1., 1., 1., 1., 1.], dtype=np.float32)

coeffs = fitter.fit_coeffs(maggies=maggies, ivar=ivar, redshift=redshift)

fitspec = coeffs.dot(templates.flux)
fitmaggies = fitter.reconstruct(redshift=redshift, coeffs=coeffs)

ab = kcorrect.utils.sed_ab(wave=templates.wave)

pngfile = 'kcorrect-lrg1-v4.fit.png'

plt.plot(np.log10(templates.wave), np.log10(fitspec / ab), color='black',
         label='Fit spectrum')
plt.scatter(np.log10(leff), np.log10(maggies), color='red', s=30, zorder=20,
            label='Observed maggies')
plt.scatter(np.log10(leff), np.log10(fitmaggies), color='black', s=20, zorder=30,
            label='Fit maggies')
plt.ylim([-1.5, 3.5])
plt.xlabel('log wavelength (Angstrom)')
plt.ylabel('log maggies')
plt.legend()
plt.savefig(pngfile, dpi=150)
