from __future__ import print_function
import os
import numpy as np
from pyemu import Matrix, Cov

"""
Explanation: Matrix and Covariance
The mat_handler.py module contains matrix class, which is the backbone of pyemu.  The matrix class overloads all common mathematical operators and also uses an "auto-align" functionality to line up matrix objects for multiplication, addition, etc.
End of explanation
"""


m = Matrix()

"""
Explanation: Here is the most basic instantiation of the matrix class:
End of explanation
"""


a = np.random.random((5, 5))
row_names = []
[row_names.append("row_{0:02d}".format(i)) for i in range(5)]
col_names = []
[col_names.append("col_{0:02d}".format(i)) for i in range(5)]
m = Matrix(x=a, row_names=row_names, col_names=col_names)
print(m)

"""
Explanation: Here we will generate a matrix object with a random ndarray
End of explanation
"""


ascii_name = "mat_test.mat"
m.to_ascii(ascii_name)
m2 = Matrix.from_ascii(ascii_name)
print(m2)

bin_name = "mat_test.bin"
m.to_binary(bin_name)
m3 = Matrix.from_binary(bin_name)
print(m3)

"""
Explanation: File I/O with matrix
matrix supports several PEST-compatible I/O routines as well as some others:
End of explanation
"""


print(type(m.to_dataframe()))
m.to_dataframe() #looks really nice in the notebook!

"""
Explanation: Matrix also implements a to_dataframe() and a to_sparse, which return pandas dataframe and a scipy.sparse (compressed sparse row) objects, respectively:
End of explanation
"""


print(m.s) #the singular values of m cast into a matrix object.  the SVD() is called on demand
m.s.to_ascii("test_sv.mat") #save the singular values to a PEST-compatible ASCII file

m.v.to_ascii("test_v.mat") #the right singular vectors of m.
m.u.to_dataframe()# a data frame of the left singular vectors of m

"""
Explanation: Convience methods of Matrix
several cool things are implemented in Matrix and accessed through @property decorated methods.  For example, the SVD components of a Matrix object are simply accessed by name.  The SVD routine is called on demand and the components are cast to Matrix objects, all opaque to the user:
End of explanation
"""


m.inv.to_dataframe()

"""
Explanation: The Matrix inverse operation is accessed the same way, but requires a square matrix:
End of explanation
"""



print(m.get(row_names="row_00",col_names=["col_01","col_03"]))

"""
Explanation: Manipulating Matrix shape
Matrix has lots of functionality to support getting submatrices by row and col names:
End of explanation
"""


from copy import deepcopy
m_copy = deepcopy(m)
sub_m = m_copy.extract(row_names="row_00",col_names=["col_01","col_03"])
m_copy.to_dataframe()
sub_m.to_dataframe()

"""
Explanation: extract() calls get() then drop():
End of explanation
"""


#a new matrix object that is not "aligned" with m
row_names = ["row_03","row_02","row_00"]
col_names = ["col_01","col_10","col_100"]
m_mix = Matrix(x=np.random.random((3,3)),row_names=row_names,col_names=col_names)
m_mix.to_dataframe()


m.to_dataframe()

prod = m * m_mix.T
prod.to_dataframe()

prod2 = m_mix.T * m
prod2.to_dataframe()

(m_mix + m).to_dataframe()

"""
Explanation: Operator overloading
The operator overloading uses the auto-align functionality as well as the isdiagonal flag for super easy linear algebra.  The "inner join" of the two objects is found and the rows and cols are aligned appropriately:
End of explanation
"""


c = Cov(m.newx,m.row_names)

"""
Explanation: The Cov derived type
The Cov type is designed specifically to handle covariance matrices.  It makes some assumptions, such as the symmetry (and accordingly that row_names == col_names).
End of explanation
"""


c.to_uncfile("test.unc")

c1 = Cov.from_uncfile("test.unc")
print(c1)

"""
Explanation: The Cov class supports several additional I/O routines, including the PEST uncertainty file (.unc):
End of explanation
"""


parcov = Cov.from_parbounds(os.path.join("henry","pest.pst"))
obscov = Cov.from_obsweights(os.path.join("henry","pest.pst"))

#to_dataframe for diagonal types builds a full matrix dataframe - can be costly
parcov.to_dataframe().head() 

# notice the zero-weight obs have been assigned a really large uncertainty
obscov.to_dataframe().head()

"""
Explanation: We can also build cov objects implied by pest control file parameter bounds or observation weights:
End of explanation
"""


