##############################################################################
# projection.py, v0.0b 2005/06/15
#
# Project:  OpenEV
# Purpose:  Projection support for OpenEV.
# Author:  Mario Beauchamp
# Comments: Very preliminary version... only works with EPSG codes for now
###############################################################################

import gtk
from gtk import TRUE,FALSE
import os
import gdal
import gvutils
import gview
import gviewapp
import osr

# those are the only drivers that work cleanly...
drivers = ['GTiff','HFA','MEM','PCIDSK','VRT']

class ProjectionTool(gviewapp.Tool_GViewApp):
    def __init__(self,app=None):
        gviewapp.Tool_GViewApp.__init__(self,app)
        self.init_menu()
        # uncomment this line if you want to create an icon entry
##        self.init_icon()

    def launch_window(self,*args):
        self.win = ProjectionDialog()
        if self.win is None:
            return
        else:
            self.win.show()

    def init_menu(self):
        self.menu_entries.set_entry("Image/Projection",1,self.launch_window)

    def init_icon(self):
        self.icon_entries.set_entry("worldrgb.xpm","Projection",1,self.launch_window)

class ProjectionDialog(gtk.GtkWindow):
    def __init__(self):
        gtk.GtkWindow.__init__(self)
        self.set_title('Projection Tool')
        self.dstSRS = osr.SpatialReference()
        self.updating = TRUE
        self.viewwin = gview.app.view_manager.get_active_view_window()
        self.active_change_id = self.viewwin.viewarea.connect("active-changed",self.activeChanged)
        self.createGUI()
        self.updating = FALSE
        self.initParams()
        self.initGUI()
        self.show_all()

    def initParams(self):
        layer = self.viewwin.viewarea.active_layer()
        if layer is None:
            return
        self.layer = layer
        self.srcDrv = layer.get_parent().get_dataset().GetDriver()
        self.dstDrv = 'GTiff'
        self.srcSRS = self.getSRS(self.layer)
        prj = self.layer.get_projection()
        if prj is not None:
            self.dstSRS.ImportFromWkt(prj)
        self.src = self.getInfoFromLayer()
        self.dst = self.compute()
        self.dst.epsg = self.src[6]

    def createGUI(self):
        self.tips = gtk.GtkTooltips()
        mainbox = gtk.GtkVBox(spacing=5)
        mainbox.set_border_width(5)
        self.add(mainbox)

# Src params
        frame = gtk.GtkFrame('Source')
        mainbox.add(frame,expand=FALSE)
        vbox = gtk.GtkVBox(spacing=5)
        vbox.set_border_width(5)
        frame.add(vbox)

        # Upper left corner
        box = gtk.GtkHBox(spacing=3)
        vbox.add(box,expand=FALSE)
        box.add(gtk.GtkLabel('UL '),expand=FALSE)
        box.add(gtk.GtkLabel('X:'),expand=FALSE)
        self.srcULxTE = GeoEntry()
        box.add(self.srcULxTE,expand=FALSE)
        box.add(gtk.GtkLabel('Y:'),expand=FALSE)
        self.srcULyTE = GeoEntry()
        box.add(self.srcULyTE,expand=FALSE)

        # Resolution
        box = gtk.GtkHBox(spacing=5)
        vbox.add(box,expand=FALSE)

        box.add(gtk.GtkLabel('Pixel size'),expand=FALSE)
        box.add(gtk.GtkLabel('X:'),expand=FALSE)
        self.srcRezxTE = gtk.GtkEntry()
        self.srcRezxTE.set_usize(100,-1)
        box.add(self.srcRezxTE,expand=FALSE)

        box.add(gtk.GtkLabel('Y:'), expand=FALSE)
        self.srcRezyTE = gtk.GtkEntry()
        self.srcRezyTE.set_usize(100,-1)
        box.add(self.srcRezyTE,expand=FALSE)

        # Size
        box = gtk.GtkHBox(spacing=5)
        vbox.add(box,expand=FALSE)

        box.add(gtk.GtkLabel('Width: '),expand=FALSE)
        self.srcWidthTE = gtk.GtkEntry()
        self.srcWidthTE.set_name('srcw')
        self.srcWidthTE.set_usize(50,-1)
        box.add(self.srcWidthTE,expand=FALSE)

        box.add(gtk.GtkLabel('Height:'),expand=FALSE)
        self.srcHeightTE = gtk.GtkEntry()
        self.srcHeightTE.set_name('srch')
        self.srcHeightTE.set_usize(50,-1)
        box.add(self.srcHeightTE,expand=FALSE)

        # SRS
        box = gtk.GtkHBox(spacing=5)
        vbox.add(box,expand=FALSE)

        box.add(gtk.GtkLabel('EPSG:'),expand=FALSE)
        self.srcEpsgTE = gtk.GtkEntry()
        self.srcEpsgTE.set_usize(50,-1)
        self.srcEpsgTE.connect('activate',self.epsgSrcChanged)
        box.add(self.srcEpsgTE,expand=FALSE)

# Dst params
        frame = gtk.GtkFrame('Destination')
        mainbox.add(frame,expand=FALSE)
        vbox = gtk.GtkVBox(spacing=5)
        vbox.set_border_width(5)
        frame.add(vbox)

        # Upper left corner
        box = gtk.GtkHBox(spacing=3)
        vbox.add(box,expand=FALSE)
        box.add(gtk.GtkLabel('UL '),expand=FALSE)
        box.add(gtk.GtkLabel('X:'),expand=FALSE)
        self.dstULxTE = GeoEntry()
        box.add(self.dstULxTE,expand=FALSE)
        box.add(gtk.GtkLabel('Y:'),expand=FALSE)
        self.dstULyTE = GeoEntry()
        box.add(self.dstULyTE,expand=FALSE)

        # Resolution
        box = gtk.GtkHBox(spacing=5)
        vbox.add(box,expand=FALSE)

        box.add(gtk.GtkLabel('Pixel size'),expand=FALSE)
        box.add(gtk.GtkLabel('X:'),expand=FALSE)
        self.dstRezxTE = gtk.GtkEntry()
        self.dstRezxTE.set_usize(100,-1)
        self.dstRezxTE.connect('activate',self.updateFromEntry,'dstx')
        box.add(self.dstRezxTE,expand=FALSE)

        box.add(gtk.GtkLabel('Y:'), expand=FALSE)
        self.dstRezyTE = gtk.GtkEntry()
        self.dstRezyTE.set_usize(100,-1)
        self.dstRezyTE.connect('activate',self.updateFromEntry,'dsty')
        box.add(self.dstRezyTE,expand=FALSE)

        # Size
        box = gtk.GtkHBox(spacing=5)
        vbox.add(box,expand=FALSE)

        box.add(gtk.GtkLabel('Width: '),expand=FALSE)
        self.dstWidthTE = gtk.GtkEntry()
        self.dstWidthTE.set_usize(50,-1)
        self.dstWidthTE.connect('activate',self.updateFromEntry,'dstw')
        box.add(self.dstWidthTE,expand=FALSE)

        box.add(gtk.GtkLabel('Height:'),expand=FALSE)
        self.dstHeightTE = gtk.GtkEntry()
        self.dstHeightTE.set_usize(50,-1)
        self.dstHeightTE.connect('activate',self.updateFromEntry,'dsth')
        box.add(self.dstHeightTE,expand=FALSE)

        # SRS
        box = gtk.GtkHBox(spacing=5)
        vbox.add(box,expand=FALSE)

        box.add(gtk.GtkLabel('EPSG:'),expand=FALSE)
        self.dstEpsgTE = gtk.GtkEntry()
        self.dstEpsgTE.set_usize(50,-1)
        self.dstEpsgTE.connect('activate',self.epsgDstChanged)
        box.add(self.dstEpsgTE,expand=FALSE)

# resampling and file create options
        frame = gtk.GtkFrame('Options')
        mainbox.add(frame,expand=FALSE)
        vbox = gtk.GtkVBox(spacing=5)
        vbox.set_border_width(5)
        frame.add(vbox)

        # warp options
        box = gtk.GtkHBox(spacing=5)
        box.set_border_width(5)
        vbox.add(box,expand=FALSE)
        box.add(gtk.GtkLabel('Resampling'),expand=FALSE)
        self.algCB = gtk.GtkCombo()
        box.add(self.algCB,expand=FALSE)
        self.algCB.set_popdown_strings(['Nearest Neighbor','Bilinear','Cubic','Cubic spline'])

        box = gtk.GtkHBox(spacing=5)
        box.set_border_width(5)
        vbox.add(box,expand=FALSE)
        box.add(gtk.GtkLabel('Error threshold'),expand=FALSE)
        self.etTE = gtk.GtkEntry()
        self.etTE.set_text('0.125')
        box.add(self.etTE,expand=FALSE)

        # file options
        box = gtk.GtkHBox(spacing=5)
        box.set_border_width(5)
        vbox.add(box,expand=FALSE)
        box.add(gtk.GtkLabel('Format'),expand=FALSE)
        self.formatCB = gtk.GtkCombo()
##        wDrvs = filter(lambda d: 'DCAP_CREATE' in d.GetMetadata(),gdal.GetDriverList())
##        drivers = map(lambda d: d.ShortName,wDrvs)
        self.formatCB.set_popdown_strings(drivers)
        self.formatCB.entry.connect('changed',self.formatChanged)
        box.add(self.formatCB,expand=FALSE)

        box = gtk.GtkHBox(spacing=5)
        box.set_border_width(5)
        vbox.add(box,expand=FALSE)
        box.add(gtk.GtkLabel('Filename'),expand=FALSE)
        self.fileTE = gtk.GtkEntry()
        box.add(self.fileTE)

        box = gtk.GtkHBox(spacing=5)
        box.set_border_width(5)
        vbox.add(box,expand=FALSE)
        box.add(gtk.GtkLabel('Create options'),expand=FALSE)
        self.coptTE = gtk.GtkEntry()
        box.add(self.coptTE)

        box = gtk.GtkHBox(homogeneous=1,spacing=5)
        box.set_border_width(5)
        vbox.add(box,expand=FALSE)
        self.viewopt0RB = gtk.GtkRadioButton(label='Current view')
        box.add(self.viewopt0RB,expand=FALSE)
        self.viewopt1RB = gtk.GtkRadioButton(label='New view',group=self.viewopt0RB)
        box.add(self.viewopt1RB,expand=FALSE)
        self.viewopt2RB = gtk.GtkRadioButton(label='Save only',group=self.viewopt0RB)
        box.add(self.viewopt2RB,expand=FALSE)
        
# Buttons
        box = gtk.GtkHBox(homogeneous=1,spacing=5)
        mainbox.add(box,expand=FALSE)

        assignBT = gtk.GtkButton("Assign")
        assignBT.connect("clicked",self.assign)
        box.add(assignBT)

        reprojBT = gtk.GtkButton("Reproject")
        reprojBT.connect("clicked",self.reproject)
        box.add(reprojBT)

        closeBT = gtk.GtkButton("Close")
        closeBT.connect("clicked",self.close)
        box.add(closeBT)

    def initGUI(self):
        self.updating = TRUE
        self.updateSrcGUI()
        self.updateDstGUI()
        self.formatCB.entry.set_text(self.dstDrv)
        self.setFilename()
        self.updating = FALSE

    def close(self,*args):
        self.viewwin.viewarea.disconnect(self.active_change_id)
        self.destroy()
        return TRUE

    def activeChanged(self,*args):
        if self.updating:
            return
        self.initParams()
        if self.layer is None:
            return
        self.initGUI()

    def epsgDstChanged(self,entry):
        if self.updating:
            return
        epsg = int(entry.get_text())
        self.dstSRS.ImportFromEPSG(epsg)
        self.dst = self.compute()
        self.dst.epsg = epsg
        self.updateDstGUI()

    def epsgSrcChanged(self,entry):
        epsg = int(entry.get_text())
        self.srcSRS.ImportFromEPSG(epsg)
        self.src[6] = epsg
        self.updateSrcGUI()

    def formatChanged(self,entry):
        self.dstDrv = entry.get_text()
        self.setFilename()
        
    def updateSrcGUI(self):
        self.setEntryMode(src=1)
        parm = map(lambda p:str(p),self.src[2:])
        self.srcULxTE.setValue(self.src[0])
        self.srcULyTE.setValue(self.src[1])
        self.srcRezxTE.set_text(parm[0])
        self.srcRezyTE.set_text(parm[1])
        self.srcWidthTE.set_text(parm[2])
        self.srcHeightTE.set_text(parm[3])
        self.srcEpsgTE.set_text(parm[4])

    def updateDstGUI(self):
        self.setEntryMode()
        self.dstULxTE.setValue(self.dst.minX)
        self.dstULyTE.setValue(self.dst.maxY)
        self.dstRezxTE.set_text(str(self.dst.rezX))
        self.dstRezyTE.set_text(str(self.dst.rezY))
        self.dstWidthTE.set_text(str(self.dst.width))
        self.dstHeightTE.set_text(str(self.dst.height))
        self.dstEpsgTE.set_text(str(self.dst.epsg))

    def updateFromEntry(self,entry,id):
        if self.updating:
            return
        self.updating = TRUE
        if id == 'dstw':
            txt = self.dstWidthTE.get_text()
            if not len(txt): return
            self.dst.setRezX(int(txt))
            self.dstRezxTE.set_text(str(self.dst.rezX))
        elif id == 'dsth':
            txt = self.dstHeightTE.get_text()
            if not len(txt): return
            self.dst.setRezY(int(txt))
            self.dstRezyTE.set_text(str(self.dst.rezY))
        elif id == 'dstx':
            txt = self.dstRezxTE.get_text()
            if not len(txt): return
            self.dst.setWidth(float(txt))
            self.dstWidthTE.set_text(str(self.dst.width))
        elif id == 'dsty':
            txt = self.dstRezyTE.get_text()
            if not len(txt): return
            self.dst.setHeight(float(txt))
            self.dstHeightTE.set_text(str(self.dst.height))

        self.updating = FALSE

    def getInfoFromLayer(self):
        params = [0,0,1,1,0,0,0]
        if self.layer is None:
            return params
        ds = self.layer.get_parent().get_dataset()
        if ds is None: # unlikely...
            return params
        params[4] = ds.RasterXSize
        params[5] = ds.RasterYSize
        geoTr = ds.GetGeoTransform()
        params[0] = geoTr[0]
        params[1] = geoTr[3]
        params[2] = geoTr[1]
        params[3] = abs(geoTr[5])
        params[6] = self.getEPSG(self.srcSRS)
        return params

    def reproject(self,but):
        Dset = self.layer.get_parent().get_dataset()

        ulx = self.dstULxTE.getValue()
        uly = self.dstULyTE.getValue()
        rezX = float(self.dstRezxTE.get_text())
        rezY = float(self.dstRezyTE.get_text())
        w = int(self.dstWidthTE.get_text())
        h = int(self.dstHeightTE.get_text())

        fname = self.fileTE.get_text()
        dstprj = self.dstSRS.ExportToWkt()
        if self.dstDrv == 'VRT': # not working... possible bug in AutoCreateWarpedVRT
            projDset = gdal.AutoCreateWarpedVRT(Dset,
                                                dst_wkt=dstprj,
                                                eResampleAlg=self.getResamplingAlg(),
                                                maxerror=float(self.etTE.get_text()))
            projDset.SetDescription(fname)
        else:
            drv = gdal.GetDriverByName(self.dstDrv)
            projDset = drv.Create(fname,w,h,
                                  bands=Dset.RasterCount,
                                  datatype=int(self.layer.type_get(0)),
                                  options=[self.coptTE.get_text()])
            if projDset is None:
                gvutils.error('Output file could not be created')
                return

        projDset.SetProjection(dstprj)
        projDset.SetGeoTransform((ulx,rezX,0.0,uly,0.0,-rezY))

        if self.dstDrv != 'VRT':
            err = gdal.ReprojectImage(Dset,projDset,
                                      dst_wkt=dstprj,
                                      eResampleAlg=self.getResamplingAlg(),
                                      maxerror=float(self.etTE.get_text()))
            if err:
                gvutils.error('Reprojection failed')
                del projDset
                return
            projDset.FlushCache()

        self.openProjected(projDset)

    def assign(self,but):
##        self.updating = TRUE
        Dset = self.layer.get_parent().get_dataset()

        ulx = self.srcULxTE.getValue()
        uly = self.srcULyTE.getValue()
        rezX = float(self.srcRezxTE.get_text())
        rezY = float(self.srcRezyTE.get_text())
        w = int(self.srcWidthTE.get_text())
        h = int(self.srcHeightTE.get_text())

##        self.viewwin.viewarea.remove_layer(self.layer)
        if self.srcDrv.ShortName == self.dstDrv:
##            projDset = Dset
##            self.layer.set_read_only(FALSE)
##            else:
                fname = Dset.GetDescription()
##                del Dset
                projDset = gdal.OpenShared(fname,gdal.GA_Update)
        else:
            drv = gdal.GetDriverByName(self.dstDrv)
            fname = self.fileTE.get_text()
            projDset = drv.CreateCopy(fname,Dset,
                                      options=[self.coptTE.get_text()])
##            del Dset

        projDset.SetProjection(self.srcSRS.ExportToWkt())
        projDset.SetGeoTransform((ulx,rezX,0.0,uly,0.0,-rezY))

##        self.updating = FALSE
        self.openProjected(projDset)

    def openProjected(self,projDset):
        if self.viewopt0RB.active:
            viewwin = self.viewwin
            viewwin.viewarea.remove_layer(self.layer)
            self.layer = None
        elif self.viewopt1RB.active:
            viewwin = self.viewwin.app.new_view()
        else:
            del projDset
            return

        viewwin.open_gdal_dataset(projDset)
        del projDset

    def compute(self):
        ulx = self.src[0]
        uly = self.src[1]
        rezX = self.src[2]
        rezY = self.src[3]
        width = self.src[4]
        height = self.src[5]
        rx = ulx+width*rezX
        ly = uly-height*rezY
        corners = [(ulx,uly),(ulx,ly),(rx,uly),(rx,ly)]

        if self.dstSRS.IsSame(self.srcSRS):
            prjCorners = corners
        else:
            ct = osr.CoordinateTransformation(self.srcSRS,self.dstSRS)
            prjCorners = ct.TransformPoints(corners)

        return GeoInfo(prjCorners,width,height)

    def getSRS(self,layer=None):
        if layer is None:
            return
        srs = osr.SpatialReference()
        prj = layer.get_projection()
        if prj is not None:
            srs.ImportFromWkt(prj)
        return srs

    def getEPSG(self,srs):
        if srs.IsGeographic():
            epsg = srs.GetAuthorityCode('GEOGCS')
        else:
            epsg = srs.GetAuthorityCode('PROJCS')
        return epsg

    def getResamplingAlg(self):
        algStr = self.algCB.entry.get_text()
        if algStr == 'Bilinear':
            alg = gdal.GRA_Bilinear
        elif algStr == 'Cubic':
            alg = gdal.GRA_Cubic
        elif algStr == 'Cubic spline':
            alg = gdal.GRA_CubicSpline
        else:
            alg = gdal.GRA_NearestNeighbour

        return alg

    def setEntryMode(self,src=0):
        if src:
            isgeo = self.srcSRS.IsGeographic()
            self.srcULxTE.setMode(isgeo)
            self.srcULyTE.setMode(isgeo)
        else:
            isgeo = self.dstSRS.IsGeographic()
            self.dstULxTE.setMode(isgeo)
            self.dstULyTE.setMode(isgeo)

    def setFilename(self):
        meta = gdal.GetDriverByName(self.dstDrv).GetMetadata()
        srcFname = os.path.splitext(self.layer.get_name())[0]
        if 'DMD_EXTENSION' in meta:
            projfn = srcFname + '-prj' + os.path.extsep + meta.get('DMD_EXTENSION')
        else:
            projfn = srcFname
        self.fileTE.set_text(projfn)

class GeoEntry(gtk.GtkHBox):
    def __init__(self,dgsz=33,mnsz=22,scsz=40,mtsz=125):
        gtk.GtkHBox.__init__(self,spacing=3)
        self.degsz = dgsz
        self.metsz = mtsz
        self.minsz = mnsz
        self.secsz = scsz
        self.round = 2

    def connect(self,signal,callback,data):
        if self.isGeo:
            self.degTE.connect(signal,callback,data)
            self.minTE.connect(signal,callback,data)
            self.secTE.connect(signal,callback,data)
        else:
            self.metersTE.connect(signal,callback,data)

    def setRound(self,value):
        self.round = value

    def setGeoMode(self):
        self.degTE = self.createEntry(self.degsz,'d')
        self.minTE = self.createEntry(self.minsz,'m')
        self.secTE = self.createEntry(self.secsz,'s')
        self.isGeo = TRUE

    def setMode(self,geo):
        if len(self.children()):
            for child in self.children():
                child.destroy()

        if geo:
            self.setGeoMode()
        else:
            self.setProjMode()

        self.show_all()

    def setProjMode(self):
        self.metersTE = self.createEntry(self.metsz,'m')
        self.isGeo = FALSE

    def createEntry(self,size,lbl):
        te = gtk.GtkEntry()
        te.set_usize(size,-1)
        self.add(te,expand=FALSE)
        self.add(gtk.GtkLabel(lbl),expand=FALSE)
        return te

    def getDegrees(self):
        txt = self.degTE.get_text()
        return self.validateInt(txt)

    def getMinutes(self):
        txt = self.minTE.get_text()
        return self.validateInt(txt)

    def getSeconds(self):
        txt = self.secTE.get_text()
        return self.validateFloat(txt)

    def getMeters(self):
        txt = self.metersTE.get_text()
        return self.validateFloat(txt)

    def validateInt(self,txt):
        if len(txt) and txt != '-':
            value = int(txt)
        else:
            value = 0

        return value

    def validateFloat(self,txt):
        if len(txt) and txt != '-':
            value = float(txt)
        else:
            value = 0.0

        return value

    def setDegrees(self,value):
        self.degTE.set_text(str(value))

    def setMinutes(self,value):
        self.minTE.set_text(str(value))

    def setSeconds(self,value):
        self.secTE.set_text(str(round(value,self.round)))

    def setMeters(self,value):
        self.metersTE.set_text(str(value))

    def setValue(self,value):
        if self.isGeo:
            dms = self.deg2dms(value)
            self.setDegrees(dms[0])
            self.setMinutes(dms[1])
            self.setSeconds(dms[2])
        else:
            self.setMeters(value)

    def set_text(self,txt):
        self.setValue(float(txt))

    def getValue(self):
        dms = []
        if self.isGeo:
            dms.append(self.getDegrees())
            dms.append(self.getMinutes())
            dms.append(self.getSeconds())
            return self.dms2deg(dms)
        else:
            return self.getMeters()

    def get_text(self):
        value = self.getValue()
        return str(value)

    def deg2dms(self,dd):
        deg = int(dd)
        dec = abs(dd-deg)
        mn = int(dec*60)
        sec = (dec*60 - mn)*60
        if round(sec,self.round) >= 60.0:
            sec = 0.0
            mn+=1
        return (deg,mn,sec)

    def dms2deg(self,dms):
        dd = abs(dms[0])+dms[1]/60.0+dms[2]/3600.0
        if dms[0] < 0:
            return -dd
        else:
            return dd

class GeoInfo:
    def __init__(self,pts,w,h):
        self.minX = min(pts[0][0],pts[1][0])
        self.maxX = max(pts[2][0],pts[3][0])
        self.minY = min(pts[1][1],pts[3][1])
        self.maxY = max(pts[0][1],pts[2][1])
        self.width = w
        self.height = h
        self.epsg = 0
        self.setRezX()
        self.setRezY()

    def setRezX(self,w=None):
        if w == None:
            w = self.width
        self.rezX = (self.maxX-self.minX)/w

    def setRezY(self,h=None):
        if h == None:
            h = self.height
        self.rezY = (self.maxY-self.minY)/h

    def setWidth(self,rez):
        self.width = int((self.maxX-self.minX)/rez)

    def setHeight(self,rez):
        self.height = int((self.maxY-self.minY)/rez)

TOOL_LIST = ['ProjectionTool']