# gdalinfo aus_for18_publish/aus_for18 -nogcp -nomd -noct -nofl | grep "<.*>"
# gdalwarp --config GDAL_CACHEMAX 10096 -multi -of GTiff -co "TILED=YES" -co "TFW=YES" -co BIGTIFF=YES -co COMPRESS=PACKBITS aus_for18 aus_for18.tiff
import xml.etree.ElementTree as ET

rat = ET.parse('rat.xml').getroot()
headers = [defn.find('Name').text for defn in rat.findall('FieldDefn')]
grouped_data = {}
id_to_label = {}
max_value = 0
id_var = 'VALUE'
grouping_var = 'FOR_TYPE'
for row in rat.findall('Row'):
    values = (dict(zip(headers, [val.text for val in row])))
    if values[grouping_var] not in grouped_data:
        grouped_data[values[grouping_var]] = []
    grouped_data[values[grouping_var]].append(values[id_var])
    id_to_label[int(values[id_var])] = values[grouping_var]
    if int(values[id_var]) > max_value:
        max_value = int(values[id_var])
#print(id_to_label)
color_table = {
    'Acacia': '#cbac25',
    'Callitris': '#2185fb',
    'Casuarina': '#722608',

    'Eucalypt Mallee Open': '#cb6768',
    'Eucalypt Mallee Woodland': '#fda07e',

    'Eucalypt Low Closed': '#9bf99c',
    'Eucalypt Low Open': '#9bf99c',
    'Eucalypt Low Woodland': '#7dc4cc',

    'Eucalypt Medium Closed': '#bdb66f',
    'Eucalypt Medium Open': '#81ce81',
    'Eucalypt Medium Woodland': '#85fd30',

    'Eucalypt Tall Closed': '#1c4306',
    'Eucalypt Tall Open': '#478915',
    'Eucalypt Tall Woodland': '#bffffe',

    'Mangrove': '#fd28fc',
    'Melaleuca': '#fffd37',
    'Rainforest': '#1a1d6f',

    'Hardwood plantation': '#ed0b19',
    'Softwood plantation': '#ed0b19',
    'Mixed species plantation': '#ed0b19',

    'Other native forest': '#fd9d2c',
    'Other forest ? unallocated type': '#eee0ce',

    'Non forest': '#ffffff',

}

from yattag import Doc, indent

doc, tag, text = Doc().tagtext()
color_indicies = list(color_table.keys())
with tag('ColorMap', type="intervals"):
    for label, quantities in grouped_data.items():
        if label not in color_table:
            print('error', label)
    for label, color in color_table.items():
        doc.stag('ColorMapEntry', color, quantity=color_indicies.index(label), label=label)

result = indent(
    doc.getvalue(),
    indentation=' ' * 4,
    newline='\r\n'
)

sld = '''
<?xml version="1.0" encoding="ISO-8859-1"?>
<StyledLayerDescriptor version="1.0.0" 
 xsi:schemaLocation="http://www.opengis.net/sld StyledLayerDescriptor.xsd" 
 xmlns="http://www.opengis.net/sld" 
 xmlns:ogc="http://www.opengis.net/ogc" 
 xmlns:xlink="http://www.w3.org/1999/xlink" 
 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
  <!-- a Named Layer is the basic building block of an SLD document -->
  <NamedLayer>
    <Name>default_raster</Name>
    <UserStyle>
    <!-- Styles can have names, titles and abstracts -->
      <Title>Default Raster</Title>
      <Abstract>A sample style that draws a raster, good for displaying imagery</Abstract>
      <!-- FeatureTypeStyles describe how to render different features -->
      <!-- A FeatureTypeStyle for rendering rasters -->
      <FeatureTypeStyle>
        <Rule>
          <RasterSymbolizer>
            <Opacity>1.0</Opacity>
            %s
          </RasterSymbolizer>
        </Rule>
      </FeatureTypeStyle>
    </UserStyle>
  </NamedLayer>
</StyledLayerDescriptor>''' % result

print(sld)
# from webcolors import hex_to_rgb
# gdaldem color-relief aus_for18 col.txt out.tif
# col = ''
# for label, quantities in grouped_data.items():
#     for quantity in quantities:
#         col += quantity + ' ' + ' '.join(str(x) for x in hex_to_rgb(color_table[label])) +'\n'
# #print(col)


import numpy
from numpy import zeros
from numpy import logical_and
from osgeo import gdal

a = numpy.array([1,2,2,1]).reshape(2,2)
# palette must be given in sorted order
palette = [-32768] + [id for (id, label) in id_to_label.items()]
# key gives the new values you wish palette to be mapped to.
key = numpy.array([-32768]+[color_indicies.index(label) for (id, label) in id_to_label.items()])
index = numpy.digitize(a.ravel(), palette, right=True)
print(key[index].reshape(a.shape))

in_file = "aus_for18_publish/aus_for18.tiff"

ds = gdal.Open(in_file)
band = ds.GetRasterBand(1)

block_sizes = band.GetBlockSize()
x_block_size = block_sizes[0]
y_block_size = block_sizes[1]

xsize = band.XSize
ysize = band.YSize

format = "GTiff"
driver = gdal.GetDriverByName( format )
dst_ds = driver.Create("original_blocks.tif", xsize, ysize, 1, gdal.GDT_Byte)#, options=['COMPRESS=PACKBITS',"TILED=YES" , "TFW=YES" ,"BIGTIFF=YES"] )
dst_ds.SetGeoTransform(ds.GetGeoTransform())
dst_ds.SetProjection(ds.GetProjection())

for i in range(0, ysize, y_block_size):
    if i + y_block_size < ysize:
        rows = y_block_size
    else:
        rows = ysize - i
    for j in range(0, xsize, x_block_size):
        if j + x_block_size < xsize:
            cols = x_block_size
        else:
            cols = xsize - j

        data = band.ReadAsArray(j, i, cols, rows)
        #r = zeros((rows, cols), numpy.uint8)
        #dst_ds.GetRasterBand(1).WriteArray(r,j,i)
        index = numpy.digitize(data.ravel(), palette, right=True)
        r = key[index].reshape(data.shape)
        dst_ds.GetRasterBand(1).WriteArray(r, j, i)
    if i % 100 == 0 or j % 100 == 0:
        print(i,j)
dst_ds = None
# gdalwarp --config GDAL_CACHEMAX 10096 -multi -of GTiff -srcnodata 0 -co "TILED=YES" -co "TFW=YES" -co BIGTIFF=YES -co COMPRESS=PACKBITS original_blocks.tif aus_for18_recolored.tiff