<commit_before>package org.apache.spark.ml.dsl.utils.messaging

import org.apache.spark.ml.dsl.utils._
import org.apache.spark.ml.util._
import org.apache.spark.sql.types.{DataType, UserDefinedType}
import org.apache.spark.util.Utils
import org.json4s.jackson.JsonMethods._
import org.json4s.{Extraction, Formats, JValue}
import org.slf4j.LoggerFactory

import scala.language.implicitConversions
import scala.util.Try
import scala.xml.{Elem, NodeSeq, XML}

//mixin to allow converting to a simple case class and back
//used to delegate ser/de tasks (from/to xml, json & dataset encoded type) to the case class with a fixed schema
//all subclasses must be objects otherwise Spark SQL can't find schema for Repr
abstract class RelayLike[Obj] {

  implicit def formats: Formats = Xml.defaultFormats

  implicit def mf: Manifest[M]

  type M //Message type

  def _read(jv: JValue, formats: Formats, mf: Manifest[M]): M = {

    val message: M = Extraction.extract[M](jv)(formats, mf)
    message
  }

  final def _fromJValue[T: MessageReader](jv: JValue): T = {

    val reader = implicitly[MessageReader[T]]
    reader._read(jv, this.formats, reader.mf)
  }
  final def _fromJSON[T: MessageReader](json: String): T = _fromJValue[T](parse(json))

  final def _fromXMLNode[T: MessageReader](ns: NodeSeq): T = {
    val jv = Xml.toJson(ns)

    _fromJValue[T](jv.children.head)
  }
  final def _fromXML[T: MessageReader](xml: String): T = {
    val nodes: Elem = xmlStr2Node(xml)

    _fromXMLNode[T](nodes)
  }
  final def _toXMLAndBack[T: MessageReader](model: T): T = {
    val xml = MessageView(model).prettyXML
    LoggerFactory.getLogger(this.getClass).info(
      s"""
         |========================= XML ========================
         |$xml
         |======================== /XML ========================
      """.stripMargin
    )
    val back = this._fromXML[T](xml)
    back
  }

  final def xmlStr2Node(xml: String): Elem = {
    val bomRemoved = xml.replaceAll("[^\\x20-\\x7e]", "").trim //remove BOM (byte order mark)
    val prologRemoved = bomRemoved.replaceFirst("[^<]*(?=<)", "")
    val result = XML.loadString(prologRemoved)
    result
  }

  implicit def reader: MessageReader[M]

  final def fromJValue(jv: JValue): M = _fromJValue[M](jv)
  final def fromJSON(json: String): M = _fromJSON[M](json)

  final def fromXMLNode(ns: NodeSeq): M = _fromXMLNode[M](ns)
  final def fromXML(xml: String): M = _fromXML[M](xml)

  def toM(v: Obj): M
  final def toMessageAPI(v: Obj): MessageAPI = {
    val m = toM(v)
    m match {
      case m: MessageAPI => m
      case _ => MessageView(m)
    }
  }
  final def toMessageAPIIfNot(v: Obj): MessageAPI = {
    v match {
      case v: MessageAPI => v
      case _ => toMessageAPI(v)
    }
  }

  trait HasMessageRelay extends MessageAPI {
    self: Obj =>

    override def formats = RelayLike.this.formats

    override def proto: Any = RelayLike.this.toM(self)
  }

  class UDT extends UserDefinedType[Obj] {

    override def sqlType: DataType = ???

    override def serialize(obj: Any): Any = ???

    override def deserialize(datum: Any): Obj = ???

    override def userClass: Class[Obj] = ???
  }

  def Param(
             parent: String,
             name: String,
             doc: String,
             isValid: Obj => Boolean,
             // serializer = SparkEnv.get.serializer
             formats: Formats = RelayLike.this.formats
           ): MessageRelayParam[Obj] = new MessageRelayParam(this, parent, name, doc, isValid, formats)

  def Param(parent: String, name: String, doc: String): MessageRelayParam[Obj] =
    Param(parent, name, doc, (_: Obj) => true)

  def Param(parent: Identifiable, name: String, doc: String, isValid: Obj => Boolean): MessageRelayParam[Obj] =
    Param(parent.uid, name, doc, isValid)

  def Param(parent: Identifiable, name: String, doc: String): MessageRelayParam[Obj] =<commit_msg>refactor and fix MessageRelay.API<commit_after>package org.apache.spark.ml.dsl.utils.messaging

import org.apache.spark.ml.dsl.utils._
import org.apache.spark.ml.util._
import org.apache.spark.sql.types.{DataType, UserDefinedType}
import org.apache.spark.util.Utils
import org.json4s.jackson.JsonMethods._
import org.json4s.{Extraction, Formats, JValue}
import org.slf4j.LoggerFactory

import scala.language.implicitConversions
import scala.util.Try
import scala.xml.{Elem, NodeSeq, XML}

//mixin to allow converting to a simple case class and back
//used to delegate ser/de tasks (from/to xml, json & dataset encoded type) to the case class with a fixed schema
//all subclasses must be objects otherwise Spark SQL can't find schema for Repr
abstract class RelayLike[Obj] {

  implicit def formats: Formats = Xml.defaultFormats

  implicit def mf: Manifest[M]

  type M //Message type

  def _read(jv: JValue, formats: Formats, mf: Manifest[M]): M = {

    val message: M = Extraction.extract[M](jv)(formats, mf)
    message
  }

  final def _fromJValue[T: MessageReader](jv: JValue): T = {

    val reader = implicitly[MessageReader[T]]
    reader._read(jv, this.formats, reader.mf)
  }
  final def _fromJSON[T: MessageReader](json: String): T = _fromJValue[T](parse(json))

  final def _fromXMLNode[T: MessageReader](ns: NodeSeq): T = {
    val jv = Xml.toJson(ns)

    _fromJValue[T](jv.children.head)
  }
  final def _fromXML[T: MessageReader](xml: String): T = {
    val nodes: Elem = xmlStr2Node(xml)

    _fromXMLNode[T](nodes)
  }
  final def _toXMLAndBack[T: MessageReader](model: T): T = {
    val xml = MessageView(model).prettyXML
    LoggerFactory.getLogger(this.getClass).info(
      s"""
         |========================= XML ========================
         |$xml
         |======================== /XML ========================
      """.stripMargin
    )
    val back = this._fromXML[T](xml)
    back
  }

  final def xmlStr2Node(xml: String): Elem = {
    val bomRemoved = xml.replaceAll("[^\\x20-\\x7e]", "").trim //remove BOM (byte order mark)
    val prologRemoved = bomRemoved.replaceFirst("[^<]*(?=<)", "")
    val result = XML.loadString(prologRemoved)
    result
  }

  implicit def reader: MessageReader[M]

  final def fromJValue(jv: JValue): M = _fromJValue[M](jv)
  final def fromJSON(json: String): M = _fromJSON[M](json)

  final def fromXMLNode(ns: NodeSeq): M = _fromXMLNode[M](ns)
  final def fromXML(xml: String): M = _fromXML[M](xml)

  def toM(v: Obj): M
  final def toMessageAPI(v: Obj): MessageAPI = {
    val m = toM(v)
    m match {
      case m: MessageAPI => m
      case _ => MessageView(m)
    }
  }
  final def toMessageAPIIfNot(v: Obj): MessageAPI = {
    v match {
      case v: MessageAPI => v
      case _ => toMessageAPI(v)
    }
  }

  trait API extends MessageAPI {
    self: Obj =>

    override def formats = RelayLike.this.formats

    final def toM = RelayLike.this.toM(self)
    final def toMessageAPI = RelayLike.this.toMessageAPI(self)
    final def toMessageAPIIfNot = RelayLike.this.toMessageAPI(self)

    override def proto: Any = toMessageAPI.proto
  }

  class UDT extends UserDefinedType[Obj] {

    override def sqlType: DataType = ???

    override def serialize(obj: Any): Any = ???

    override def deserialize(datum: Any): Obj = ???

    override def userClass: Class[Obj] = ???
  }

  def Param(
             parent: String,
             name: String,
             doc: String,
             isValid: Obj => Boolean,
             // serializer = SparkEnv.get.serializer
             formats: Formats = RelayLike.this.formats
           ): MessageRelayParam[Obj] = new MessageRelayParam(this, parent, name, doc, isValid, formats)

  def Param(parent: String, name: String, doc: String): MessageRelayParam[Obj] =
    Param(parent, name, doc, (_: Obj) => true)

  def Param(parent: Identifiable, name: String, doc: String, isValid: Obj => Boolean): MessageRelayParam[Obj] =
    Param(parent.uid, name, doc, isValid)

  def Param(parent: Identifiable, name: String, doc: String): MessageRelayParam[Obj] =
{'fraction_non_alphanumeric': 0.1450783699059561, 'fraction_numerical': 0.002006269592476489, 'mean_word_length': 3.4212860310421287, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 26}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6892386', 'n_tokens_mistral': 2706, 'n_tokens_neox': 2619, 'n_words': 691}