#!python

"""
This script contains tsx_quad pol reader
and UnitTest
"""

from tsx_dual import tsx_dual,latlon2xy
import unittest
import osgeo.gdal as gdal
import numpy as np

class tsx_quad(tsx_dual):

    def __init__(self,hh_pol="",vv_pol="",hv_pol="",vh_pol="",slc=""):
        self.slc_ds = gdal.Open(slc)
        self.meta_xml = slc
        self.slc_transform = gdal.GCPsToGeoTransform(self.slc_ds.GetGCPs())
        self.xsize = self.slc_ds.RasterXSize
        self.ysize = self.slc_ds.RasterYSize

    def calc_raw_amp_list(self,point_list,win):
        """
        Calculates raw amplitude
        e.g. all points in a field
        """

        slc_transform = self.slc_transform

        hh_list = list()
        hv_list = list()
        vv_list = list()
        vh_list = list()

        #DEBUG
        #logfile = file("C:\\test_xy.txt","a")

        for point in point_list:
            slc_x,slc_y = latlon2xy(point[0],point[1],slc_transform)

            #DEBUG
            #logfile.write(str(slc_x)+","+str(slc_y)+",")
            try:
                hh = self.slc_ds.GetRasterBand(1).ReadAsArray(slc_x,slc_y,win,win)
                hv = self.slc_ds.GetRasterBand(2).ReadAsArray(slc_x,slc_y,win,win)
                vh = self.slc_ds.GetRasterBand(3).ReadAsArray(slc_x,slc_y,win,win)
                vv = self.slc_ds.GetRasterBand(4).ReadAsArray(slc_x,slc_y,win,win)
            except:
                hh = np.zeros((win,win))
                hv = np.zeros((win,win))
                vh = np.zeros((win,win))
                vv = np.zeros((win,win))

            hh_amp = np.abs(hh)
            hv_amp = np.abs(hv)
            vh_amp = np.abs(vh)
            vv_amp = np.abs(vv)

            hh_list.append(np.mean(hh_amp))
            hv_list.append(np.mean(hv_amp))
            vh_list.append(np.mean(vh_amp))
            vv_list.append(np.mean(vv_amp))

        #DEBUG
        #logfile.close()

        return hh_list,hv_list,vh_list,vv_list

    def in_data(self,latlon):
        slc_transform = self.slc_transform
        x,y = latlon2xy(latlon[0],latlon[1],slc_transform)
        return ((x > 0 and y > 0) and (x < self.xsize and y < self.ysize))

    def read_area(self,tlbr):
        tllat = tlbr[0][0]
        tllon = tlbr[0][1]
        brlat = tlbr[1][0]
        brlon = tlbr[1][1]

        xstart,ystart = latlon2xy(tllat,tllon,self.slc_transform)
        xend,yend = latlon2xy(brlat,brlon,self.slc_transform)

        tlcornerx = min(xstart,xend)
        tlcornery = min(ystart,yend)
        xsize = abs(xend-xstart)
        ysize = abs(yend-ystart)

        try:
            return self.slc_ds.ReadAsArray(tlcornerx,tlcornery,
                                           xsize,ysize)
        except:
            return None


class TestTsxQuad(unittest.TestCase):
    def setUp(self):
        tsx_file = "C:\\imagery_temp\dims_op_oc_dfd2_205662493_1\\"+\
        "TSX-1.SAR.L1B\\TSX1_SAR__SSC______SM_Q_DRA_20100419T201037"+\
        "_20100419T201045\\TSX1_SAR__SSC______SM_Q_DRA_20100419T201037"+\
        "_20100419T201045.xml"
        self.test_loc = ((-34.534190,138.748411),)
        self.fail_loc = ((0.0,0,0),)
        self.test_area = ((-34.534190,138.748411),(-34.6,138.8))
        self.tsx_data = tsx_quad(slc=tsx_file)

    def test_valid_gdal(self):
        """
        Valid GDAL dataset has been initialized
        """
        self.assertNotEqual(self.tsx_data.slc_ds,None)

    def test_band_count(self):
        """
        Quad pol data has 4-bands
        """
        self.assertEqual(self.tsx_data.slc_ds.RasterCount,4)

    def test_size(self):
        """
        Data size is non-zero
        """
        self.assertNotEqual(self.tsx_data.slc_ds.RasterXSize,0)
        self.assertNotEqual(self.tsx_data.slc_ds.RasterYSize,0)

    def test_read(self):
        """
        Data can be read from GDAL
        """
        data = self.tsx_data.slc_ds.ReadAsArray(0,0,1,1)
        self.assertEqual(data.shape[0],4)

    def test_read_outside(self):
        """
        Read outside extents is 0
        """
        mult_data = self.tsx_data.calc_raw_amp_list(self.fail_loc, 1)
        self.assertNotEqual(mult_data,None)
        self.assertEqual(len(mult_data),4)
        self.assertEqual(len(mult_data[0]),1)
        for channel in mult_data:
            self.assertEqual(channel[0],0)

    def test_read_multiple(self):
        """
        Read in extents is non-zero [town]
        """
        mult_data = self.tsx_data.calc_raw_amp_list(self.test_loc, 1)
        self.assertNotEqual(mult_data,None)
        self.assertEqual(len(mult_data),4)
        self.assertEqual(len(mult_data[0]),1)
        for channel in mult_data:
            self.assertNotEqual(channel[0],0)

    def test_read_area(self):
        """
        Read a set of pixels given topleft and bottom right
        """
        mult_data = self.tsx_data.read_area(self.test_area)
        self.assertNotEqual(mult_data,None)
        self.assertNotEqual(np.mean(mult_data),0)

    def test_in_data(self):
        """
        Do latlon in data bounds checking
        """
        self.assertTrue(self.tsx_data.in_data(self.test_loc[0]))
        self.assertFalse(self.tsx_data.in_data(self.fail_loc[0]))


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