import os
import warnings
from rasterio import rasterio, shutil
from pyproj.transformer import Transformer
from pyproj.crs import CRS
from geopy.point import Point
from PIL import Image
import numpy as np

class GeoTiff:
    def __init__(self, path):
        with warnings.catch_warnings(record=True) as w: # отлавливаем ошибки
            self.path = path
            self.file = rasterio.open(path)
            if len(w) > 0 and issubclass(w[-1].category, rasterio.errors.NotGeoreferencedWarning): # если файл имеет контрольные точки
                self.file = rasterio.vrt.WarpedVRT(self.file, src_crs=self.file.gcps[1], scrs=self.file.gcps[1])    # приводим его к виду georeferenced


    # получить размеры изображения
    def get_size(self):
        return self.file.width, self.file.height


    def get_crs(self):
        return self.file.crs


    def get_transform(self):
        return self.file.transform


    # получить тип данных пикселя
    def get_dtype(self):
        return self.file.read(1).dtype


    # получить координаты крайних точек изображения
    def get_corner_coordinates(self):
        width, height = self.get_size()
        return [
            self._transform_to_coordinates(*self.file.xy(0, 0)),
            self._transform_to_coordinates(*self.file.xy(0, width)),
            self._transform_to_coordinates(*self.file.xy(height, 0)),
            self._transform_to_coordinates(*self.file.xy(height, width)),
        ]


    # получить координаты по индексу
    def get_coordinate_by_index(self, width, height):
        row, col = self.file.xy(height, width)
        return self._transform_to_coordinates(row, col)


    # получить индексы координаты
    def get_index_by_coordinate(self, coordinate):
        x, y = self._transform_to_meters(coordinate)
        return self.file.index(x, y)


    # получить numpy отрезок снимка по индексам
    def get_map_by_indexes(self, width1, height1, width2, height2, channel=1):
        return self.file.read(channel)[height1:height2, width1:width2]


    # установить numpy отрезок снимка по индексам TODO check
    def set_map_by_indexes(self, width, height, slice_map, channel=1):
        channel_map = self.file.read(channel)
        channel_map[height1:height2, width1:width2] = slice_map
        return self.file.write(channel_map, channel)


    # получить numpy отрезок снимка по координатам
    def get_map_by_coordinates(self, coordinate1, coordinate2, channel=1):
        width1, height1 = self.get_index_by_coordinate(coordinate1)
        width2, height2 = self.get_index_by_coordinate(coordinate2)
        return self.get_map_by_indexes(width1, height1, width2, height2, channel)


    # изменить значение прозрачности
    def change_transparency(self, transparency=255, pixel_value=0):
        if self.file.count < 4:
            return
        width, height = self.get_size()
        tr = np.zeros((height, width), dtype=self.get_dtype())
        for band in range(1, self.file.count):
            image_map = self.get_map_by_indexes(0, 0, width, height, band)
            for h in range(height):
                for w in range(width):
                    if image_map[h][w] == pixel_value:
                        tr[h][w] = transparency      
        self.file.write(tr, 4)


    # сохранить часть изображения по индексам
    def save_image_by_indexes(self, path, width1, height1, width2, height2, shape=False):
        width, height = self.get_size()
        if width1 > width2 or height1 > height2 or width1 > width or height1 > height:
            return False
        if width2 > width:
            if shape:
                return False
            width2 = width
        if height2 > height:
            if shape:
                return False
            height2 = height
        if self.file.count == 1:
            mode = "L"
            image_map = self.get_map_by_indexes(width1, height1, width2, height2)
        elif self.file.count == 3:
            mode = "RGB"
            image_map = np.stack((
                np.reshape(self.get_map_by_indexes(width1, height1, width2, height2, 1), (height2-height1, width2-width1, 1)), 
                np.reshape(self.get_map_by_indexes(width1, height1, width2, height2, 2), (height2-height1, width2-width1, 1)),
                np.reshape(self.get_map_by_indexes(width1, height1, width2, height2, 3), (height2-height1, width2-width1, 1)),
            ), axis=2)
        else:
            mode = "RGBA"
            image_map = np.stack((
                np.reshape(self.get_map_by_indexes(width1, height1, width2, height2, 1), (height2-height1, width2-width1, 1)), 
                np.reshape(self.get_map_by_indexes(width1, height1, width2, height2, 2), (height2-height1, width2-width1, 1)),
                np.reshape(self.get_map_by_indexes(width1, height1, width2, height2, 3), (height2-height1, width2-width1, 1)),
                np.reshape(self.get_map_by_indexes(width1, height1, width2, height2, 4), (height2-height1, width2-width1, 1)),
            ), axis=2)
        Image.fromarray(image_map, mode=mode).save(path)
        return True


    # сохранить часть изображения по координатам
    def save_image_by_coordinates(self, path, coordinate1, coordinate2, mode="L"):
        width1, height1 = self.get_index_by_coordinate(coordinate1)
        width2, height2 = self.get_index_by_coordinate(coordinate2)
        return self.save_image_by_indexes(width1, height1, width2, height2, mode)


    # сохранить снимок в виде Georeferenced (по умолчанию перезаписать)
    def save_file_as_georeferenced(self, path=None):
        if path is None:
            path = self.path
        shutil.copy(self.file, path, driver='GTiff')
        return True


    # перевести метрическую систему координат в географическую
    def _transform_to_coordinates(self, x, y):
        transformer = Transformer.from_proj(self.file.crs, CRS("EPSG:4326"))
        lat, lon = transformer.transform(x, y)
        lon = abs(lon+90)
        if abs(lon-90) > abs(lon-180):
            if abs(lon-180) > abs(lon-270):
                lon = abs(270-lon) + 90
            else:
                lon = lon - 90
        return Point(lat, lon)


    # перевести географическую систему координат в метрическую
    def _transform_to_meters(self, coordinate):
        transformer = Transformer.from_crs(CRS("EPSG:4326"), self.file.crs)
        x, y = transformer.transform(coordinate.latitude, coordinate.longitude)
        return y, x
