from __future__ import annotations
"""Object writing helpers."""

import os
from typing import Any

import numpy as np

import vedo
import vedo.vtkclasses as vtki
from vedo import utils
from vedo.assembly import Assembly
from vedo.mesh import Mesh

from .loaders import load
from .scene import to_numpy

__docformat__ = "google"

def write(objct: Any, fileoutput: str | os.PathLike, binary=True) -> Any:
    """
    Write object to file. Same as `save()`.

    Supported extensions are:

    - `vtk, vti, ply, obj, stl, byu, vtp, vti, mhd, xyz, xml, tif, png, bmp`
    """
    fileoutput = str(fileoutput)

    ###############################
    if isinstance(objct, Assembly):
        dd = to_numpy(objct)
        sdict = {"objects": [dd]}
        np.save(fileoutput, sdict)
        return objct

    ###############################
    obj = objct.dataset

    # Check if object actor has a non-identity transform and bake it before writing.
    try:
        M = objct.actor.GetMatrix()
    except AttributeError:
        M = None
    if M and not M.IsIdentity():
        objct.apply_transform_from_actor()
        obj = objct.dataset
        vedo.logger.info(
            f"object '{objct.name}' "
            "was manually moved. Writing uses current position."
        )

    fr = fileoutput.lower()
    if fr.endswith(".vtk"):
        writer = vtki.new("DataSetWriter")
    elif fr.endswith(".ply"):
        writer = vtki.new("PLYWriter")
        writer.AddComment("PLY file generated by vedo")
        lut = objct.mapper.GetLookupTable()
        if lut:
            pscal = obj.GetPointData().GetScalars()
            if not pscal:
                pscal = obj.GetCellData().GetScalars()
            if pscal and pscal.GetName():
                writer.SetArrayName(pscal.GetName())
            writer.SetLookupTable(lut)
    elif fr.endswith(".stl"):
        writer = vtki.new("STLWriter")
    elif fr.endswith(".vtp"):
        writer = vtki.new("XMLPolyDataWriter")
    elif fr.endswith(".vtu"):
        writer = vtki.new("XMLUnstructuredGridWriter")
    elif fr.endswith(".xyz"):
        writer = vtki.new("SimplePointsWriter")
    elif fr.endswith(".facet"):
        writer = vtki.new("FacetWriter")
    elif fr.endswith(".vti"):
        writer = vtki.new("XMLImageDataWriter")
    elif fr.endswith(".vtr"):
        writer = vtki.new("XMLRectilinearGridWriter")
    elif fr.endswith(".vtm"):
        g = vtki.new("MultiBlockDataGroupFilter")
        for ob in objct:
            try:
                g.AddInputData(ob)
            except TypeError:
                vedo.logger.warning(f"cannot save object of type {type(ob)}")
        g.Update()
        mb = g.GetOutputDataObject(0)
        wri = vtki.new("vtkXMLMultiBlockDataWriter")
        wri.SetInputData(mb)
        wri.SetFileName(fileoutput)
        wri.Write()
        return objct
    elif fr.endswith(".mhd"):
        writer = vtki.new("MetaImageWriter")
    elif fr.endswith(".nii"):
        writer = vtki.new("NIFTIImageWriter")
    elif fr.endswith(".png"):
        writer = vtki.new("PNGWriter")
    elif fr.endswith(".jpg"):
        writer = vtki.new("JPEGWriter")
    elif fr.endswith(".bmp"):
        writer = vtki.new("BMPWriter")
    elif fr.endswith(".tif") or fr.endswith(".tiff"):
        writer = vtki.new("TIFFWriter")
        writer.SetFileDimensionality(len(obj.GetDimensions()))
    elif fr.endswith(".obj"):
        with open(fileoutput, "w", encoding="UTF-8") as outF:
            outF.write("# OBJ file format with ext .obj\n")
            outF.write("# File generated by vedo\n")

            for p in objct.vertices:
                outF.write("v {:.8g} {:.8g} {:.8g}\n".format(*p))

            ptxt = objct.dataset.GetPointData().GetTCoords()
            if ptxt:
                ntxt = utils.vtk2numpy(ptxt)
                for vt in ntxt:
                    outF.write("vt " + str(vt[0]) + " " + str(vt[1]) + " 0.0\n")

            if isinstance(objct, Mesh):
                for i, f in enumerate(objct.cells):
                    fs = ""
                    for fi in f:
                        if ptxt:
                            fs += f" {fi+1}/{fi+1}"
                        else:
                            fs += f" {fi+1}"
                    outF.write(f"f{fs}\n")

                for l in objct.lines:
                    ls = ""
                    for li in l:
                        ls += str(li + 1) + " "
                    outF.write(f"l {ls}\n")
        return objct

    elif fr.endswith(".off"):
        with open(fileoutput, "w", encoding="UTF-8") as outF:
            outF.write("OFF\n")
            outF.write(str(objct.npoints) + " " + str(objct.ncells) + " 0\n\n")
            for p in objct.vertices:
                outF.write(" ".join([str(i) for i in p]) + "\n")
            for c in objct.cells:
                outF.write(str(len(c)) + " " + " ".join([str(i) for i in c]) + "\n")
        return objct

    else:
        vedo.logger.error(f"Unknown format {fileoutput}, file not saved")
        return objct

    try:
        if binary:
            writer.SetFileTypeToBinary()
        else:
            writer.SetFileTypeToASCII()
    except AttributeError:
        pass

    try:
        writer.SetInputData(obj)
        writer.SetFileName(fileoutput)
        writer.Write()
    except Exception as e:
        vedo.logger.error(f"could not save {fileoutput}: {e}")
    return objct

def save(obj: Any, fileoutput="out.png", binary=True) -> Any:
    """Save an object to file. Same as `write()`."""
    return write(obj, fileoutput, binary)

def read(obj: Any, unpack=True, force=False) -> Any:
    """Read an object from file. Same as `load()`."""
    return load(obj, unpack, force)


__all__ = ["write", "save", "read"]
