#!/usr/bin/env python

# Copyright (C) 2026 Smithsonian Astrophysical Observatory
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#

'''
Create a region file to visualize cross match results

    green are from the input source list
    magenta are from the reference source list

    filled in regions are included in the final xform solution

    a line will be drawn between all cross match pairs.
'''


__toolname__ = "xmatch_viz"
__revision__ = "09 May 2026"

import sys
import ciao_contrib.logger_wrapper as lw
from pycrates import read_file


def log_wrapper(func):
    'wrapper around logger to check for None'
    def wrapped(msg):
        if msg:
            func(msg)
    return wrapped


lw.initialize_logger(__toolname__)
verb0 = log_wrapper(lw.get_logger(__toolname__).verbose0)
verb1 = log_wrapper(lw.get_logger(__toolname__).verbose1)
verb2 = log_wrapper(lw.get_logger(__toolname__).verbose2)


class Source():
    'everything about a source region'

    color = "green"
    shape = "circle"
    tag = "Input"

    def __init__(self, crate, index):
        """
        Initialize a Source object with the given crate and index.

        Parameters
        ----------
        crate : pycrates.Crate
            The crate object containing source data
        index : int
            The row index (1-based) of the source in the crate

        Attributes
        ----------
        ra : float
            Right ascension of the source
        dec : float
            Declination of the source
        matched : bool
            Whether the source has been matched (ref_index > 0)
        include : bool
            Whether the source is included in the analysis
        ref_index : int
            Reference index for matched sources (-999 if unmatched)
        """
        self.crate = crate
        self.index = index

        self.ra = None
        self.dec = None
        self.matched = None
        self.include = None
        self.ref_index = -999
        self.load()

    def load(self):
        """Load data from row in crate.

        Populates source attributes (ra, dec, include, ref_index, matched)
        by reading the corresponding row from the crate table.
        """
        _idx = self.index - 1  # row number to python array index

        self.ra = self.crate.get_column("RA").values[_idx]
        self.dec = self.crate.get_column("DEC").values[_idx]
        self.include = self.crate.get_column("INCLUDE").values[_idx]
        self.ref_index = self.crate.get_column("REF_INDEX").values[_idx]
        self.matched = self.ref_index > 0

    def get_shape(self, radius):
        'Construct the shape string'
        shape = f"{self.shape}({self.ra},{self.dec},{radius}\")"
        return shape

    def render(self, pars):
        """Create a region string for the source.

        Returns a formatted region description for the source with a circle
        shape at the source coordinates, the fill state based on inclusion,
        and tags indicating the source type and match state.
        """

        shape = self.get_shape(float(pars["radius"]))
        fill = 1 if self.include else 0

        state = "Unmatched"
        if self.matched:
            state = "Matched"
            if fill:
                state = "Used"

        src = f"{shape} # text={{{self.index}}} fill={fill} "
        src += f"color={self.color} tag={{All_{self.tag}}} tag={{{self.tag}_{state}}}\n"
        return src

    def link_cross_match(self, ref_src_list):
        """Create a line linking this source to its matched reference source.

        If this source has a valid reference index, construct a line region
        string between the source and the corresponding reference source.
        Also mark the referenced source as matched and propagate the inclusion
        state.

        Parameters
        ----------
        ref_src_list : ReferenceSourceList
            The list of reference sources containing the matched source.

        Returns
        -------
        str
            A region string describing the line between the source and the
            matched reference source, or an empty string if there is no
            match.
        """

        line = ""
        if ref_src_list is None:
            return line

        ref_idx = self.ref_index-1
        if self.ref_index > 0:
            ref_ra = ref_src_list.srcs[ref_idx].ra
            ref_dec = ref_src_list.srcs[ref_idx].dec

            shape = f"line({self.ra},{self.dec},{ref_ra},{ref_dec})"
            color = "white"
            dash = 0 if self.include else 1
            line = f"{shape} # line=0 0 color={color} dash={dash}\n"

            ref_src_list.srcs[ref_idx].matched = True
            ref_src_list.srcs[ref_idx].include = self.include

        return line


class ReferenceSource(Source):
    'Sub class for reference sources'

    color = "magenta"
    tag = "Reference"
    shape = "box"

    def load(self):
        """Load only the right ascension and declination for this reference source.

        The reference source is represented by a crate and an index. This method
        reads the RA and DEC columns from the crate and stores the values on the
        instance.
        """
        _idx = self.index - 1
        self.ra = self.crate.get_column("RA").values[_idx]
        self.dec = self.crate.get_column("DEC").values[_idx]

    def get_shape(self, radius):
        'Create the shape string'
        shape = f'{self.shape}({self.ra},{self.dec},{2*radius}",{2*radius}",0)'
        return shape


class SourceList():
    'a list of sources'

    def __init__(self, infile):
        """Initialize a SourceList with the given input file.

        Args:
            infile: Path to the input file containing source data.
        """
        self.infile = infile
        self.num_src = 0
        self.srcs = []
        self.load()

    def load(self):
        """Load the table from the input file and populate sources.

        Reads the source data from the input file, determines the number of
        sources, and creates Source objects for each row in the table.
        """
        tab = read_file(self.infile)
        self.num_src = tab.get_nrows()
        for idx in range(self.num_src):
            self.srcs.append(Source(tab, idx+1))

    def render(self, pars, ref_src_list=None):
        """Create regions for all sources and generate cross-match links.

        Args:
            ref_src_list: Optional reference source list for cross-matching.
                         Defaults to None.

        Returns:
            tuple: A tuple containing:
                - out_src (str): Concatenated region output from all sources.
                - out_line (str): Concatenated cross-match link output from all sources.
        """
        out_src = ""
        out_line = ""
        for src in self.srcs:
            out_src += src.render(pars)
            out_line += src.link_cross_match(ref_src_list)

        return (out_src, out_line)


class ReferenceSourceList(SourceList):
    'List of references sources'

    def load(self):
        """Load reference sources from input file.

        Reads the input file and loads reference sources containing only
        their RA/DEC coordinates. Populates the srcs list with ReferenceSource
        objects and sets the total number of sources.
        """
        tab = read_file(self.infile)
        self.num_src = tab.get_nrows()
        for idx in range(self.num_src):
            self.srcs.append(ReferenceSource(tab, idx+1))


class Director():
    'Write to output file'

    def __init__(self, outfile, legend, clobber):
        """Initialize a Director for writing output files.

        Args:
            outfile (str): Path to the output file.
            clobber (bool): Whether to overwrite existing files.
        """

        from ciao_contrib._tools.fileio import outfile_clobber_checks
        self.outfile = outfile
        self.legend = legend
        self.__fp = None
        outfile_clobber_checks(clobber, self.outfile)
        outfile_clobber_checks(clobber, self.legend)

    def __write_header(self):
        """Write DS9 region file header information.

        Writes the DS9 version header, global region properties (color, line style,
        font, etc.), and coordinate frame specification (FK5) to the output file.
        """

        # Note: edit and move are disabled.
        self.__fp.write("# Region file format: DS9 version 4.1\n")
        self.__fp.write('global color=green dashlist=8 3 width=2 ')
        self.__fp.write('font="helvetica 10 normal roman" select=1 ')
        self.__fp.write('highlite=1 dash=0 fixed=0 edit=0 move=0 delete=1 include=1 source=1\n')
        self.__fp.write("fk5\n")

    def write_region(self, src, ref, lines):
        """Write a DS9 region file with the specified regions and lines.

        Args:
            src (str): String containing source regions to write.
            ref (str): String containing reference regions to write.
            lines (str): Additional lines to include in the region file.
        """

        with open(self.outfile, "w", encoding="ascii") as self.__fp:
            self.__write_header()
            self.__fp.write(lines)
            self.__fp.write(src)
            self.__fp.write(ref)

    def write_legend(self):
        """Write a fixed DS9 illustrate-mode legend file.

        The legend file contains a fixed set of DS9 shapes and labels
        illustrating the meaning of the various source and match symbols.
        """
        with open(self.legend, "w", encoding="ascii") as fp:
            fp.write('# Illustrate file format: DS9 version 1.0\n')
            fp.write('global color = cyan fill = no width = 1 dash = no\n')
            fp.write('global font = helvetica fontsize = 12 fontweight = normal fontslant = roman\n')
            fp.write('circle 55.0 45.0 15.0  # color = forestgreen width = 2.0\n')
            fp.write('box 55.0 90.0 15.0 15.0 # color = magenta width = 2.0\n')
            fp.write('circle 55.0 135.0 15.0  # color = forestgreen fill = yes width = 2.0\n')
            fp.write('box 55.0 180.0 15.0 15.0  # color = magenta fill = yes width = 2.0\n')
            fp.write('circle 55.0 225.0 15.0  # color = forestgreen width = 2.0\n')
            fp.write('box 105.0 225.0 15.0 15.0  # color = magenta width = 2.0\n')
            fp.write('circle 55.0 270.0 15.0  # color = forestgreen fill = yes width = 2.0\n')
            fp.write('box 105.0 270.0 15.0 15.0  # color = magenta fill = yes width = 2.0\n')
            fp.write('text 157.0 45.0 "Unused Input Source" # color = white\n')
            fp.write('text 176.0 90.0 "Unused Reference Source" # color = white\n')
            fp.write('text 149.0 135.0 "Used Input Source" # color = white\n')
            fp.write('text 170.0 180.0 "Used Reference Source" # color = white\n')
            fp.write('text 177.0 225.0 "Unused match" # color = white\n')
            fp.write('text 168.0 270.0 "Used match" # color = white\n')
            fp.write('box 80.0 270.0 25.0 2.0  # color = white fill = yes\n')
            fp.write('box 80.0 225.0 25.0 2.0  # color = white fill = yes\n')


def postface(pars):
    "Provide ds9 command line to load files"

    verb1(" ")
    verb1("Use the following command to load the regions into DS9:")
    verb1(" ")
    verb1("ds9 FILE_NAME \\")
    verb1(f"    -region format ds9 -region {pars["outfile"]} \\")
    verb1(f"    -illustrate load {pars["legend"]} \\")
    verb1(f"    -catalog import fits {pars["infile"]}")
    verb1(" ")
    verb1("replacing FILE_NAME with the name of your event file or image.")


@lw.handle_ciao_errors(__toolname__, __revision__)
def main():
    'Main routine'
    # Load parameters
    from ciao_contrib.param_soaker import get_params
    pars = get_params(__toolname__, "rw", sys.argv,
                      verbose={"set": lw.set_verbosity, "cmd": verb1},
                      revision=__revision__)

    src_obj = SourceList(pars["infile"])
    ref_obj = ReferenceSourceList(pars["refsrcfile"])
    outfile = Director(pars["outfile"], pars["legend"], pars["clobber"])

    src_reg, line_reg = src_obj.render(pars, ref_obj)
    ref_reg, _ = ref_obj.render(pars)

    outfile.write_region(src_reg, ref_reg, line_reg)
    outfile.write_legend()
    postface(pars)


def test():
    'Test routine'
    bleck = SourceList("ngc4755_21169.xmatch[crossmatch]")
    ref = ReferenceSourceList("gaia_dr3.fits")

    goo = Director("bleck.reg", "bleck.seg", True)
    ss, ll = bleck.render({'radius': 0.5}, ref)
    rr, _ = ref.render({'radius': 0.5})
    goo.write_region(ss, rr, ll)


if __name__ == '__main__':
    main()
