#!/usr/bin/env python

# coding:utf-8

#

# Description:

#

#

# Author: LC

# Date:

#

try:

    from osgeo import gdal

except ImportError:

    import gdal

import Queue

import threading

import time

import os

import osr

import tarfile

import math

exitFlag = 0

# 自定义多线程类

class myThread(threading.Thread):

    def __init__(self, threadID, name, queueLock, workQueue, outputdir, demfile):

        threading.Thread.__init__(self)

        self.threadID = threadID

        self.name = name

        self.queueLock = queueLock

        self.workQueue = workQueue

        self.outputdir = outputdir

        self.demfile = demfile

    def run(self):

        print "Starting " + self.name

        GF4_Ortho(self.name, self.queueLock, self.workQueue, self.outputdir, self.demfile)

        print "Exiting " + self.name

def GF4_Ortho(threadName, queueLock, workQueue, outputdir, demfile):

    while not exitFlag:

        queueLock.acquire()

        if not workQueue.empty():

            untarfile = workQueue.get()

            queueLock.release()

            print "%s processing %s" % (threadName, untarfile)

            # 解压

            t = tarfile.open(untarfile)

            outdir_ = outputdir + os.path.sep + ((os.path.basename(untarfile)).split(".tar"))[0]

            t.extractall(path = outdir_)

            # 计算UTM区号

            filename = ((os.path.basename(untarfile)).split('.tar'))[0]

            lon = float(((filename.split('_'))[2])[1:])

            zone_ = int(math.ceil(lon / 6)) + 30

            zone = int("" + str(zone_))

            print zone

            # 正射校正

            filename_ = filename.split('_')

            filename_[1] = 'PMS'

            infile = outdir_ + os.path.sep + '_'.join(filename_) + '.tiff'

            outfile = outdir_ + os.path.sep + '_'.join(filename_) + '_Ortho.tiff'

            if os.path.exists(infile):

                dataset = gdal.Open(infile)

                dstSRS = osr.SpatialReference()

                dstSRS.ImportFromEPSG(zone)

                ret = gdal.Warp(outfile, dataset, format="Gtiff", xRes=50, yRes=50, dstSRS=dstSRS, rpc=True,

                                transformerOptions=demfile)

                del t, infile, outfile, dataset, ret

        else:

            queueLock.release()

        time.sleep(1)

def main(inputdir, outputdir, demfile, thread_number):

    # 查找待处理压缩文件

    names = os.listdir(inputdir)

    filenames = []

    for name_ in names:

        filenames.append(inputdir + os.path.sep + name_)

    num = len(filenames)

    threadList = []

    for i in range(thread_number):

        threadList.append("Thread-" + str(i))

    queueLock = threading.Lock()

    workQueue = Queue.Queue(num)

    threads = []

    threadID = 1

    # 创建新线程

    for tName in threadList:

        thread = myThread(threadID, tName, queueLock, workQueue, outputdir, demfile)

        thread.start()

        threads.append(thread)

        threadID += 1

    # 填充队列

    queueLock.acquire()

    for name in filenames:

        workQueue.put(name)

    queueLock.release()

    # 等待队列清空

    while not workQueue.empty():

        pass

    # 通知线程是时候退出

    exitFlag = 1

    # 等待所有线程完成

    for t in threads:

        t.join()

    print "Exiting Main Thread"

if __name__ == "__main__":

    indir = "H:\\GF-4\\2016-05-22\\"

    outdir = "G:\\Temp\\GF-4\\"

    demfile = "C:\\Program Files\\Exelis\\ENVI53\data\\GMTED2010.jp2"

    thread_num = 6

    starttime = time.ctime()

    main(indir, outdir, demfile, thread_num)

    endtime = time.ctime()

    print "Start time:", starttime

    print "End time:", endtime