# Ladybug: A Plugin for Environmental Analysis (GPL)
# This file is part of Ladybug.
#
# Copyright (c) 2025, Ladybug Tools.
# You should have received a copy of the GNU Affero General Public License
# along with Ladybug; If not, see <http://www.gnu.org/licenses/>.
# 
# @license AGPL-3.0-or-later <https://spdx.org/licenses/AGPL-3.0-or-later>

"""
Color a mesh as a heatmap using values that align with the mesh faces or vertices.
_
Note that any brep can be converted to a gridded mesh that can be consumed by 
this component using the "LB Generate Point Grid" component.
-

    Args:
        _values: A list of numerical values with which to color the mesh.
            The number of values must match the number of faces or vertices
            in the mesh.
        _mesh: A Mesh object, with a number of faces or vertices that match
            the number of input values and will be colored with results.
        offset_dom_: Optional domain (or number for distance), which will
            be used to offset the mesh faces or verticesto according to the
            values. Higher values will be offset further.
        legend_par_: Optional legend parameters from the Ladybug
            'Legend Parameters' component.
        legend_title_: A text string for Legend title. Typically, the units
            of the data are used here but the type of data might also be used.
            Default is an empty string.
        global_title_: A text string to label the entire mesh.  It will be
            displayed in the lower left of the result mesh. Default is for no
            title.

    Returns:
        mesh: The input _mesh that has been colored with results.
        legend: Geometry representing the legend for the mesh.
        title: A text object for the global_title.
        colors: The colors associated with each input value.
        legend_par: The input legend parameters with defaults filled for
            unset properties.
        vis_set: A VisualizationSet object for drawing a detailed version of the
            heatmap in the Rhino scene. This can be connected to
            the "LB Preview Visualization Set" component to display this
            version of the heatmap in Rhino.
"""

ghenv.Component.Name = 'LB Spatial Heatmap'
ghenv.Component.NickName = 'Heatmap'
ghenv.Component.Message = '1.9.1'
ghenv.Component.Category = 'Ladybug'
ghenv.Component.SubCategory = '4 :: Extra'
ghenv.Component.AdditionalHelpFromDocStrings = '1'

try:
    from ladybug.legend import LegendParameters
except ImportError as e:
    raise ImportError('\nFailed to import ladybug:\n\t{}'.format(e))

try:
    from ladybug_display.visualization import VisualizationSet
except ImportError as e:
    raise ImportError('\nFailed to import ladybug:\n\t{}'.format(e))

try:
    from ladybug_rhino.togeometry import to_mesh3d
    from ladybug_rhino.fromgeometry import from_mesh3d
    from ladybug_rhino.fromobjects import legend_objects
    from ladybug_rhino.text import text_objects
    from ladybug_rhino.color import color_to_color
    from ladybug_rhino.grasshopper import all_required_inputs
except ImportError as e:
    raise ImportError('\nFailed to import ladybug_rhino:\n\t{}'.format(e))


if all_required_inputs(ghenv.Component):
    # translate to Ladybug objects
    lb_mesh = to_mesh3d(_mesh)
    if offset_dom_:
        dom_st, dom_end = offset_dom_
        lb_mesh = lb_mesh.height_field_mesh(_values, (dom_st, dom_end))

    # check the values against the mesh
    assert len(_values) == len(lb_mesh.faces) or len(_values) == len(lb_mesh.vertices), \
        'Expected the number of data set values ({}) to align with the number of faces ' \
        '({}) or the number of vertices ({}).\nConsider flattening the _values input ' \
        'and using the "Mesh Join" component to join the _mesh input.'.format(
            len(_values), len(lb_mesh.faces), len(lb_mesh.vertices))

    # create the VisualizationSet and GraphicContainer
    if legend_title_ is not None:
        legend_par_ = legend_par_.duplicate() if legend_par_ is not None \
            else LegendParameters()
        legend_par_.title = legend_title_
    vis_set = VisualizationSet.from_single_analysis_geo(
        'Data_Mesh', [lb_mesh], _values, legend_par_)
    graphic = vis_set.graphic_container()

    # generate titles
    if global_title_ is not None:
        title = text_objects(global_title_, graphic.lower_title_location,
                             graphic.legend_parameters.text_height * 1.5,
                             graphic.legend_parameters.font)

    # draw rhino objects
    lb_mesh.colors = graphic.value_colors
    mesh = from_mesh3d(lb_mesh)
    legend = legend_objects(graphic.legend)
    colors = [color_to_color(col) for col in lb_mesh.colors]
    legend_par = graphic.legend_parameters
