#!/usr/bin/env python

import gdal
import gdalconst
import sys
import Numeric
import os.path
import osr
import colorsys
import Image
import glob
import string
import os
import pickle

def mask_min(masked_band, mask, x, y):
	(size_y, size_x) = Numeric.shape(mask)
	current = masked_band.ReadAsArray(x, y, size_x, size_y)
	masked_band.WriteArray(Numeric.minimum(current, mask), x, y)

def find_x_pos(ct_to_lonlat, geomatrix, y, target_lon, hint_x=0):
	(lon, lat) = transform_pixel(ct_to_lonlat, geomatrix, (hint_x, y))

	x = hint_x

	# If we're to the right of the left edge, move left until we're
	# left of the edge and then add one.
	if lon > target_lon:
		while (lon > target_lon):
			x -= 1
			(lon, lat) = transform_pixel(ct_to_lonlat, geomatrix, (x, y))
		return(x + 1, lat)
		
	# If we're to the left of the left edge, move right until we're
	# right of the edge.
	else:
		while (lon < target_lon):
			x += 1
			(lon, lat) = transform_pixel(ct_to_lonlat, geomatrix, (x, y))
		return(x, lat)

def find_y_pos(ct_to_lonlat, geomatrix, x, target_lat, hint_y=0):
	(lon, lat) = transform_pixel(ct_to_lonlat, geomatrix, (x, hint_y))

	y = hint_y

	# If we're above the bottom edge, move down until we're
	# below the edge and then subtract one.
	if lat > target_lat:
		while (lat > target_lat):
			y += 1
			(lon, lat) = transform_pixel(ct_to_lonlat, geomatrix, (x, y))
		return(y - 1, lon)
		
	# If we're below the bottom edge, move up until we're
	# above the edge.
	else:
		while (lat < target_lat):
			y -= 1
			(lon, lat) = transform_pixel(ct_to_lonlat, geomatrix, (x, y))
		return(y, lon)

def transform_pixel(ct, geomatrix, pixel):
	(pixel_x, pixel_y) = pixel
	X = geomatrix[0] + geomatrix[1] * pixel_x + geomatrix[2] * pixel_y
	Y = geomatrix[3] + geomatrix[4] * pixel_x + geomatrix[5] * pixel_y
	(lon, lat, height) = ct.TransformPoint(X, Y)
	return(lon, lat)

def transform_lonlat(ct, geomatrix, lonlat):
	(a, b, c, d, e, f) = geomatrix
	(lon, lat) = lonlat
	(X, Y, Z) = ct.TransformPoint(lon, lat)
	x = int((c * (Y - d) - f * (X - a)) / (c * e - b * f))
	y = int((b * (Y - d) - e * (X - a)) / (b * f - c * e))
	return(x, y)

if 1:
	sectionals_path = '/var/www/aviationtoolbox/raw_data/FAA_sectionals/current/'
	masked_sectionals_path = './masked/'
	warped_sectionals_path = './warped/'

	# Load info about the center of our universe.
	reference_filename = glob.glob('../raw_data/FAA_sectionals/current/Wichita*North.tif')[0]
	reference_image = gdal.Open(reference_filename, gdalconst.GA_ReadOnly)
	reference_projection = reference_image.GetProjection()
	reference_srs = osr.SpatialReference()
	reference_srs.ImportFromWkt(reference_projection)
	reference_image = None

	horizontal_skip_min = 850
	horizontal_skip_max = 2000
	right_border = 30
	bottom_border = 30
	warped_x_pixel_size = 64.0
	warped_y_pixel_size = 64.0

	masked_bands = warped_bands = 4

	sectionals_info_file = open('./sectional_corners', 'r')
	
	for sectional_info_line in sectionals_info_file.readlines():
		print sectional_info_line
		(sectional_region, ll_corner_string, notes) = string.split(sectional_info_line, '|')
		#if not sectional_region == 'Chicago':
		#	continue

		ll_corner_lonlat = tuple(map(float, string.split(ll_corner_string, ' ')))

		sectional_version = glob.glob('%s%s ?? North.tif' % (sectionals_path, sectional_region))[0][-12:-10]

		# for sectional_side in ['North', 'South']:
		for sectional_side in ['South']:
			sectional_filename = '%s%s %s %s.tif' % (sectionals_path, sectional_region, sectional_version, sectional_side)
			masked_filename = '%s%s %s.tif' % (masked_sectionals_path, sectional_region, sectional_side)
			warped_filename = '%s%s %s.tif' % (warped_sectionals_path, sectional_region, sectional_side)
			info_filename = 'masked/%s %s.info' % (sectional_region, sectional_side)

		
			if 0 and os.path.exists(masked_filename):
				print 'already have %s, skipping' % (masked_filename)
				continue

			sectional_image = gdal.Open(sectional_filename, gdalconst.GA_ReadOnly)
			sectional_band = sectional_image.GetRasterBand(1)
			sectional_color_table = sectional_band.GetRasterColorTable()
			sectional_geomatrix = sectional_image.GetGeoTransform()
			sectional_srs = osr.SpatialReference()
			sectional_projection = sectional_image.GetProjection()
			sectional_srs.ImportFromWkt(sectional_projection)
			srs_lonlat = sectional_srs.CloneGeogCS()
			sectional_ct_to_lonlat = osr.CoordinateTransformation(sectional_srs, srs_lonlat)
			sectional_ct_to_pixel = osr.CoordinateTransformation(srs_lonlat, sectional_srs)
			sectional_to_reference_ct = osr.CoordinateTransformation(sectional_srs, reference_srs)

			masked_image = gdal.GetDriverByName('GTiff').Create(
				masked_filename,
				sectional_image.RasterXSize, sectional_image.RasterYSize,
				masked_bands, gdal.GDT_Byte, # R,G,B,alpha,mask
				[ 'TILED=YES' ]
			)
			masked_image.SetGeoTransform(sectional_geomatrix)
			masked_image.SetProjection(sectional_projection)
			masked_band = masked_image.GetRasterBand(1)
			masked_band.SetDescription(sectional_band.GetDescription())
			masked_band.SetMetadata(sectional_band.GetMetadata())
			# masked_band.SetRasterColorTable(sectional_band.GetRasterColorTable())
#
			colorish_index = Numeric.ones (256) * 4

			masked_band_r = masked_image.GetRasterBand(1)
			masked_band_g = masked_image.GetRasterBand(2)
			masked_band_b = masked_image.GetRasterBand(3)
			#masked_band_a = masked_image.GetRasterBand(4)
			masked_band_m = masked_image.GetRasterBand(4) # mask quality
		
			# This doesn't seem to work.
			#masked_band_r.SetNoDataValue(255)
			#masked_band_g.SetNoDataValue(255)
			#masked_band_b.SetNoDataValue(255)
			#masked_band_a.SetNoDataValue(0)
			#masked_band_m.SetNoDataValue(255)
			# Initialize the image manually.
			blank_lines = Numeric.ones((sectional_image.RasterYSize, sectional_image.RasterXSize)) * 255
			
			# For testing.
			# masked_band_m = masked_band_r

			#masked_band_r.WriteArray(blank_lines, 0, 0)
			#masked_band_g.WriteArray(blank_lines, 0, 0)
			#masked_band_b.WriteArray(blank_lines, 0, 0)
			masked_band_m.WriteArray(blank_lines, 0, 0)

			# Create the palette.
			lookup = [
				Numeric.arrayrange(257),
				Numeric.arrayrange(257),
				Numeric.arrayrange(257),
				Numeric.ones(257)*255
			]

			for i in range (min (256, sectional_color_table.GetCount ())):
				#print i,
				(r, g, b, x) = sectional_color_table.GetColorEntry (i)
				(h, l, s) = colorsys.rgb_to_hls(r, g, b)
				#print '%0.3f, %0.3f, %03d' % (h, l, s), 

				if s > -0.04 and l > 128:
					colorish_index[i] = 0
				else:
					colorish_index[i] = 1

				# If it's completely white, pull a dash of blue.
				#if r + g + b == 3 * 255:
				#	b = 254

				lookup[0][i] = r	
				lookup[1][i] = g	
				lookup[2][i] = b

			# I think this is a mask.
			colorish_index[255] = -1

			# Where is the lower left corner?
			ll_corner_pixel = transform_lonlat(sectional_ct_to_pixel, sectional_geomatrix, ll_corner_lonlat)

			# We'll use these a lot.
			(ll_corner_x, ll_corner_y) = ll_corner_pixel
			(ll_corner_lon, ll_corner_lat) = ll_corner_lonlat

			x_min = x_max = sectional_image.RasterXSize
			y_min = y_max = 0

			# Find the top line.
			inmap_vertical_count = 0
			for iY in range(0, sectional_image.RasterYSize, 1):
				sectional_data = sectional_band.ReadAsArray(
					ll_corner_x, iY,
					sectional_image.RasterXSize - ll_corner_x, 1
				)
				sectional_data_indexed = Numeric.take(colorish_index, sectional_data)[0]
						
				colorval = sum(sectional_data_indexed) * 100 / len(sectional_data_indexed)

				# print colorval

				if colorval >= 80: 
					inmap_vertical_count += 1
					if inmap_vertical_count > 3:
						y_min = iY - 3
						break
				else:
					inmap_vertical_count = 0


			# Mask out everything above top line as 0.
			masked_band_m.WriteArray(Numeric.zeros((y_min, sectional_image.RasterXSize)), 0, 0)

			# Feather mask below the top line.
			mask = Numeric.fromfunction(lambda y, x: y, (255, sectional_image.RasterXSize))
			# It's the first one, so don't bother checking what's there.
			masked_band_m.WriteArray(mask, 0, y_min)
			
			# Feather mask the right edge.
			mask = Numeric.fromfunction(lambda y, x: 255 - x, (sectional_image.RasterYSize, 255))
			mask_min(masked_band_m, mask, sectional_image.RasterXSize - 255, 0)

			# Start at the x position of the LL corner just to get close.
			left_edge_x = ll_corner_x

			mask = Numeric.fromfunction(lambda y, x: x, (1, 255))
			# Find the left edge.
			for iY in range(y_min, sectional_image.RasterYSize, 1):
				# Use the previous left_edge_x as a starting point.
				(left_edge_x, lat) = find_x_pos(sectional_ct_to_lonlat, sectional_geomatrix, iY, ll_corner_lon, hint_x=left_edge_x)

				# If we're below the LL corner, stop.
				if lat < ll_corner_lat:
					break

				x_min = min(left_edge_x, x_min)

				# Mask out to the left of the left edge.
				masked_band_m.WriteArray(Numeric.zeros((1, left_edge_x)), 0, iY)

				# Feather mask to the right of the left_edge.
				#masked_band_m.WriteArray(mask, left_edge_x, iY)
				mask_min(masked_band_m, mask, left_edge_x, iY)

			# If we reached the bottom of the image, feather mask above bottom edge.
			if lat >= ll_corner_lat:
				y_max = sectional_image.RasterYSize
				mask = Numeric.fromfunction(lambda y, x: 255 - y, (255, sectional_image.RasterXSize - left_edge_x))
				mask_min(masked_band_m, mask, left_edge_x, sectional_image.RasterYSize - 255)
	
			# Otherwise we are below the LL corner now.
			else:
				# Mask out left and below the LL corner.
				masked_band_m.WriteArray(Numeric.zeros((sectional_image.RasterYSize - ll_corner_y, ll_corner_x)), 0, ll_corner_y)

				# Follow the bottom edge to the right.
				bottom_edge_y = ll_corner_y
				for iX in range(ll_corner_x, sectional_image.RasterXSize):
					(bottom_edge_y, lon) = find_y_pos(sectional_ct_to_lonlat, sectional_geomatrix, iX, ll_corner_lat, hint_y=bottom_edge_y)

					y_max = max(bottom_edge_y, y_max)

					# Mask out below.
					masked_band_m.WriteArray(Numeric.zeros((sectional_image.RasterYSize - bottom_edge_y, 1)), iX, bottom_edge_y)
					# Feather mask above.
					mask = Numeric.fromfunction(lambda y, x: 255 - y, (255, 1))
					mask_min(masked_band_m, mask, iX, bottom_edge_y - 255)
					
			# Feather mask the right edge.
			mask = Numeric.fromfunction(lambda y, x: 255 - x, (sectional_image.RasterYSize, 255))
			mask_min(masked_band_m, mask, sectional_image.RasterXSize - 255, 0)

			# Copy over the sectional image, making completely masked (0) pixels white.
			# I suppose this could be done all at once instead of line by line?
			gdal.TermProgress(0.0)
			for iY in range(sectional_image.RasterYSize):
			#for iY in range(30):
				mask = masked_band_m.ReadAsArray(0, iY, sectional_image.RasterXSize, 1)

				sectional_band_data = sectional_band.ReadAsArray(0, iY, sectional_image.RasterXSize, 1)
				for band in range(3):
					band_lookup = lookup[band]
					masked_data = Numeric.array([Numeric.ones(sectional_image.RasterXSize) * 255])
					Numeric.putmask(masked_data, mask, Numeric.take(band_lookup, sectional_band_data))
					masked_image.GetRasterBand(band + 1).WriteArray(masked_data, 0, iY)


				gdal.TermProgress( (iY+1.0) / sectional_image.RasterYSize )

			masked_image = None




			break





			corners = [
				(x_min, y_min),
				(x_min, y_max),
				(x_max, y_min),
				(x_max, y_max),
			]

			targ_xmin = None
			for corner in corners:
				(x, y) = corner
				X = sectional_geomatrix[0] + sectional_geomatrix[1] * x + sectional_geomatrix[2] * y
				Y = sectional_geomatrix[3] + sectional_geomatrix[4] * x + sectional_geomatrix[5] * y
				(x, y, z) = sectional_to_reference_ct.TransformPoint(X, Y)
				print corner, x,y
				if targ_xmin is None:
					targ_xmin = targ_xmax = x
					targ_ymin = targ_ymax = y
				else:
					targ_xmin = min(targ_xmin, x)
					targ_xmax = max(targ_xmax, x)
					targ_ymin = min(targ_ymin, y)
					targ_ymax = max(targ_ymax, y)

			warped_xsize = (targ_xmax - targ_xmin) / warped_x_pixel_size + 1
			warped_ysize = (targ_ymax - targ_ymin) / warped_y_pixel_size + 1
			print warped_xsize, warped_ysize

			warped_gt = (targ_xmin, warped_x_pixel_size, 0.0, targ_ymax, 0.0, - warped_y_pixel_size)

			warped_image = gdal.GetDriverByName('GTiff').Create(
				warped_filename,
				warped_xsize, warped_ysize,
				warped_bands, gdal.GDT_Byte, # RGBM
				[ 'TILED=YES' ]
			)

			warped_image.SetProjection( reference_srs.ExportToWkt() )
			warped_image.SetGeoTransform( warped_gt )
			# warped_image.GetRasterBand(1).SetRasterColorTable(sectional_image.GetRasterBand(1).GetRasterColorTable())
			#warped_image.GetRasterBand(1)
			#warped_image.GetRasterBand(2)
			#warped_image.GetRasterBand(3)
			warped_image.GetRasterBand(1).SetNoDataValue(255)
			warped_image.GetRasterBand(2).SetNoDataValue(255)
			warped_image.GetRasterBand(3).SetNoDataValue(255)

			blank_lines = Numeric.ones((warped_ysize, warped_xsize)) * 255
			warped_image.GetRasterBand(1).WriteArray(blank_lines,0,0)
			warped_image.GetRasterBand(2).WriteArray(blank_lines,0,0)
			warped_image.GetRasterBand(3).WriteArray(blank_lines,0,0)
			# Should I initialize mask band?

			warped_image = None

			#command = './warp "%s" "%s"' % (masked_filename, warped_filename)
			#print command
			#os.system(command)
			#command = 'convert -scale 10%% "%s" "%s.png"' % (warped_filename, warped_filename[:-4])
			#print command
			#os.system(command)

			break
		break